diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index cbce61e5e..32a4cb00d 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -8,17 +8,27 @@ on: 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] concurrency: group: ${{ github.workflow }}-${{ 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 +# uses a dedicated PAT secret), so this is behavior-neutral. Jobs that need +# more must opt in at the job level. +permissions: + contents: read + env: - PRIMUS_TURBO_COMMIT: 333b68d7c81b722b21b4aad10cd250c45f15027c # fix sm_scale none bug (#263) - PRIMUS_TURBO_AITER_COMMIT: e83f9903c07001a0ec29e85d223f6e6cdbe00859 + PRIMUS_TURBO_COMMIT: 350ec3f594a35c835a3f9ec53e97285f7bd8e7a7 # dev/kyle/flydsl_attn_deepseekv4 (dsv4 sparse-MLA attention + cr=4 fixes) + 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) - BASE_IMAGE: docker.io/rocm/primus:v26.2 + TRITON_COMMIT: 88b227e23f0445f3f695bad05bbf1a363b4f50e0 + BASE_IMAGE: docker.io/rocm/primus:v26.3 MAXTEXT_BASE_IMAGE: docker.io/rocm/jax-training:maxtext-v26.2 jobs: @@ -30,10 +40,10 @@ jobs: steps: - run: echo "🎉 Begin Primus Python Lint." - run: echo "🐧 This job is now running on a ${{ runner.os }} server hosted by GitHub!" - - uses: actions/checkout@v4 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 - run: git config --global --add safe.directory /github/workspace - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 with: python-version: ${{ matrix.python-version }} - name: Display Python version @@ -44,15 +54,32 @@ jobs: - name: Run pre-commit (all hooks, all files) run: pre-commit run --all-files --show-diff-on-failure + # 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. + dependency-review: + if: github.event_name == 'pull_request' + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - name: Dependency Review + uses: actions/dependency-review-action@a1d282b36b6f3519aa1f3fc636f609c47dddb294 # v5.0.0 + with: + fail-on-severity: high + warn-only: true + 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 strategy: matrix: python-version: ["3.12"] steps: - run: echo "🎉 Begin Build Primus Docker Image." - - uses: actions/checkout@v4 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 with: submodules: recursive - name: Show Environment Info @@ -93,6 +120,7 @@ jobs: --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 end_time=$(date +%s) elapsed=$((end_time - start_time)) @@ -103,25 +131,25 @@ jobs: docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}} # docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }} - # # Primus v26.2 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-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.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 - # end_time=$(date +%s) - # elapsed=$((end_time - start_time)) - # echo "⏱️ [build primus docker-ainic] Total elapsed time: ${elapsed} seconds" + # Primus v26.3 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 \ + --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 + end_time=$(date +%s) + elapsed=$((end_time - start_time)) + echo "⏱️ [build primus docker-ainic] Total elapsed time: ${elapsed} seconds" - # docker tag tasimage/primus:${{env.IMAGE_TAG}}-ainic docker.io/tasimage/primus:${{env.IMAGE_TAG}}-ainic - # docker login -u tasimage -p ${{ secrets.PRIMUS_DOCKER_HUB_TOKEN }} - # docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}-ainic + docker tag tasimage/primus:${{env.IMAGE_TAG}}-ainic docker.io/tasimage/primus:${{env.IMAGE_TAG}}-ainic + docker login -u tasimage -p ${{ secrets.PRIMUS_DOCKER_HUB_TOKEN }} + 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" @@ -153,7 +181,8 @@ jobs: --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 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" @@ -190,218 +219,479 @@ jobs: # docker rmi tasimage/primus:${{env.IMAGE_TAG}}-jax echo "> build-docker success" - run-unittest-torch: - env: - PRIMUS_WORKDIR: /mnt/apps_proxy/tas/0_public/primus_ci/actions-runner-torch - # PRIMUS_WORKDIR: /wekafs/primus-data/primus_safe_ci/torch - needs: [code-lint] - # runs-on: [primus-lm-cicd-torch-j8knc] - runs-on: [primus-lm-cicd-v26.2-tas8n-a16-40] - steps: - - run: echo "🎉 Begin Primus-Turbo Checkout." - - name: Set commit hash to env - run: echo "PRIMUS_TURBO_COMMIT=${PRIMUS_TURBO_COMMIT}" >> $GITHUB_ENV - - name: Checkout Repo Primus-Turbo - uses: actions/checkout@v4 - with: - repository: AMD-AIG-AIMA/Primus-Turbo - submodules: "recursive" - path: Primus-Turbo - ref: ${{ env.PRIMUS_TURBO_COMMIT }} - - run: echo "Begin AITER + Primus-Turbo Install." - - name: Install AITER - run: | - echo "✅ [Uninstall old aiter] started at: $(date)" - pip3 uninstall aiter amd-aiter -y || true - rm -rf /tmp/aiter || true - cd /tmp - git clone --recursive https://github.com/ROCm/aiter.git - cd aiter - git checkout ${PRIMUS_TURBO_AITER_COMMIT} - 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" - - 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" - 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" - - run: echo "🎉 Begin Primus Unit Test." - - uses: actions/checkout@v4 - with: - submodules: recursive - - name: Show Environment Info - run: | - echo "Hostname: $(hostname)" - echo "PWD: $(pwd)" - echo "HOME: $HOME" - echo "GITHUB_WORKSPACE: $GITHUB_WORKSPACE" - echo "Runner Temp Dir: $RUNNER_TEMP" - echo "Runner Tool Cache: $RUNNER_TOOL_CACHE" - - name: Install Primus - run: | - pip install -r requirements.txt - - name: Set UT_LOG_PATH - run: | - ts="$(date +%Y%m%d-%H%M%S)" - commit_id="${GITHUB_SHA::7}" - if [[ "${{ github.event_name }}" == "pull_request" ]]; then - 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 }}" == "release" ]]; then - TAG_NAME="${{ github.ref }}" - TAG="${TAG_NAME#refs/tags/}" - echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/${TAG}-${ts}-${commit_id}" >> $GITHUB_ENV - else - echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/others-${ts}-${commit_id}" >> $GITHUB_ENV - fi - - name: Run CLI Shell Tests - run: | - echo "Running Primus CLI shell tests..." - bash ./tests/runner/run_all_tests.sh - - 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 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 - pytest --maxfail=1 -s ./tests/unit_tests/ \ - --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 - - name: Run Primus Model Tests -- Megatron-LM - env: - HF_TOKEN: ${{secrets.HF_TOKEN}} - run: | - 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_megatron_trainer.py - - name: Run Primus Model Tests -- TorchTitan - env: - HF_TOKEN: ${{secrets.HF_TOKEN}} - run: | - 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 - - name: Clean - run: | - rm -rf ${PRIMUS_WORKDIR}/Primus-Turbo - rm -rf ${PRIMUS_WORKDIR}/Primus +# run-unittest-torch: +# env: +# 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] +# 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] +# 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." +# - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 +# with: +# submodules: recursive +# - name: Show Environment Info +# run: | +# echo "Hostname: $(hostname)" +# echo "PWD: $(pwd)" +# echo "HOME: $HOME" +# echo "GITHUB_WORKSPACE: $GITHUB_WORKSPACE" +# echo "Runner Temp Dir: $RUNNER_TEMP" +# echo "Runner Tool Cache: $RUNNER_TOOL_CACHE" +# - name: Install Primus +# run: | +# 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 +# - name: Set UT_LOG_PATH +# run: | +# ts="$(date +%Y%m%d-%H%M%S)" +# commit_id="${GITHUB_SHA::7}" +# if [[ "${{ github.event_name }}" == "pull_request" ]]; then +# 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 }}" == "release" ]]; then +# TAG_NAME="${{ github.ref }}" +# TAG="${TAG_NAME#refs/tags/}" +# echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/${TAG}-${ts}-${commit_id}" >> $GITHUB_ENV +# else +# echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/others-${ts}-${commit_id}" >> $GITHUB_ENV +# fi +# - name: Run CLI Shell Tests +# run: | +# echo "Running Primus CLI shell tests..." +# bash ./tests/runner/run_all_tests.sh +# - 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 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 +# 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} \ +# --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 +# - name: Snapshot unit-test coverage +# if: always() +# continue-on-error: true +# run: | +# # 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 +# - 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. +# # parallel=true -> one data file per rank; sigterm=true captures ranks +# # killed on teardown. These env vars are inherited by the E2E steps. +# 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" +# - 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. +# M=1; T=1 +# 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)" +# echo "Selected torch E2E scope: ${e2e:-}" +# if [[ "${e2e}" != "all" ]]; then +# echo "${e2e}" | grep -qw megatron || M=0 +# echo "${e2e}" | grep -qw torchtitan || T=0 +# fi +# fi +# fi +# echo "RUN_MEGATRON_E2E=${M}" >> "$GITHUB_ENV" +# echo "RUN_TORCHTITAN_E2E=${T}" >> "$GITHUB_ENV" +# - 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 \ +# --junitxml="${GITHUB_WORKSPACE}/test-reports/megatron-e2e.xml" +# - 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 \ +# --junitxml="${GITHUB_WORKSPACE}/test-reports/torchtitan-e2e.xml" +# - name: Run Primus CLI tool smoke (benchmark / preflight, E2E coverage) +# if: always() +# continue-on-error: true +# run: | +# # 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" +# 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" +# cli benchmark gemm-dense --seqlen 256 --hidden-size 512 --intermediate-size 512 --duration 2 --output-file "$OUT/gemm_dense.md" +# cli benchmark gemm-deepseek --seqlen 256 --hidden-size 512 --intermediate-size 512 --moe-intermediate-size 512 --duration 2 --output-file "$OUT/gemm_deepseek.md" +# cli benchmark attention --backend flash --mbs-list 1,2 --report-csv-path "$OUT/attn.csv" +# cli preflight --host --gpu --dump-path "$OUT/preflight" +# # multi-GPU collectives (best-effort; strided needs >=8 GPUs) +# 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 +# - name: Upload test reports +# if: always() +# uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 +# with: +# name: torch-test-reports +# path: test-reports/ +# if-no-files-found: warn +# retention-days: 14 +# - name: Write test summary +# if: always() +# run: | +# python tools/ci/junit_summary.py --title torch test-reports/*.xml >> "$GITHUB_STEP_SUMMARY" || true +# - name: Write runtime summary +# if: always() +# run: | +# python tools/ci/runtime_summary.py --title torch "$RUNNER_TEMP/runtime.tsv" >> "$GITHUB_STEP_SUMMARY" || true +# - name: Build torch coverage json (unit + E2E) +# if: always() +# continue-on-error: true +# run: | +# # 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 +# - name: Upload torch coverage json +# if: always() +# continue-on-error: true +# uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 +# with: +# name: coverage-torch +# path: | +# coverage_unit.json +# coverage_combined.json +# if-no-files-found: ignore +# - 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 - run-unittest-jax: - 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] - runs-on: [primus-jax-tas-runner] # docker container primus_jax_github_runner on tas a16-31 - steps: - - run: echo "🎉 Begin Primus-Turbo Checkout." - - name: Set commit hash to env - run: echo "PRIMUS_TURBO_COMMIT=${PRIMUS_TURBO_COMMIT}" >> $GITHUB_ENV - - name: Checkout Repo Primus-Turbo - uses: actions/checkout@v4 - with: - repository: AMD-AIG-AIMA/Primus-Turbo - submodules: "recursive" - path: Primus-Turbo - ref: ${{ env.PRIMUS_TURBO_COMMIT }} - - run: echo "Begin Primus-Turbo Install." - - 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 - cd /tmp/Primus-Turbo - start_time=$(date +%s) - echo "✅ [Pip install requirements] started at: $(date)" - mkdir -p ${PRIMUS_WORKDIR}/primus-cache - python3 -m pip install --upgrade pip setuptools - 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" - start_time=$(date +%s) - echo "✅ [build primus-turbo] started at: $(date)" - end_time=$(date +%s) - elapsed=$((end_time - start_time)) - echo "✅ [build primus-turbo] ended at: $(date)" - echo "⏱️ [build primus-turbo] Torch installation causes segfault, so we skip it and actually not install turbo. Total elapsed time: ${elapsed} seconds" - - run: echo "🎉 Begin Primus Unit Test." - - uses: actions/checkout@v4 - with: - submodules: recursive - - name: Show Environment Info - run: | - echo "Hostname: $(hostname)" - echo "PWD: $(pwd)" - echo "HOME: $HOME" - echo "GITHUB_WORKSPACE: $GITHUB_WORKSPACE" - echo "Runner Temp Dir: $RUNNER_TEMP" - echo "Runner Tool Cache: $RUNNER_TOOL_CACHE" - - name: Install Primus - run: | - pip install -r requirements-jax.txt - - name: Set UT_LOG_PATH - run: | - ts="$(date +%Y%m%d-%H%M%S)" - commit_id="${GITHUB_SHA::7}" - if [[ "${{ github.event_name }}" == "pull_request" ]]; then - 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 }}" == "release" ]]; then - TAG_NAME="${{ github.ref }}" - TAG="${TAG_NAME#refs/tags/}" - echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/${TAG}-${ts}-${commit_id}" >> $GITHUB_ENV - else - echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/others-${ts}-${commit_id}" >> $GITHUB_ENV - fi - - name: Run Shell Tests - run: | - echo "Running Primus CLI shell tests..." - bash ./tests/runner/run_all_tests.sh - - name: Run Unit Tests - env: - HF_TOKEN: ${{secrets.HF_TOKEN}} - run: | - 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 \ - JAX_SKIP_UT=1 python ./tests/run_unit_tests.py --jax - - name: Clean - run: | - rm -rf ${PRIMUS_WORKDIR}/Primus-Turbo - rm -rf ${PRIMUS_WORKDIR}/Primus +# run-unittest-jax: +# 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." +# - 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 +# path: Primus-Turbo +# ref: ${{ env.PRIMUS_TURBO_COMMIT }} +# - name: Init Primus-Turbo submodules (full clone) +# working-directory: Primus-Turbo +# run: | +# rm -rf 3rdparty/composable_kernel .git/modules/3rdparty/composable_kernel +# git submodule sync --recursive +# git submodule update --init --recursive +# - run: echo "Begin Primus-Turbo Install." +# - 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 +# cd /tmp/Primus-Turbo +# : > "$RUNNER_TEMP/runtime.tsv" # reset the CI runtime log for this job +# start_time=$(date +%s) +# echo "✅ [Pip install requirements] started at: $(date)" +# mkdir -p ${PRIMUS_WORKDIR}/primus-cache +# python3 -m pip install --upgrade pip setuptools +# 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 "pip install/upgrade\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv" +# start_time=$(date +%s) +# echo "✅ [build primus-turbo] started at: $(date)" +# end_time=$(date +%s) +# elapsed=$((end_time - start_time)) +# echo "✅ [build primus-turbo] ended at: $(date)" +# echo "⏱️ [build primus-turbo] Torch installation causes segfault, so we skip it and actually not install turbo. Total elapsed time: ${elapsed} seconds" +# - run: echo "🎉 Begin Primus Unit Test." +# - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 +# with: +# submodules: recursive +# - name: Show Environment Info +# run: | +# echo "Hostname: $(hostname)" +# echo "PWD: $(pwd)" +# echo "HOME: $HOME" +# echo "GITHUB_WORKSPACE: $GITHUB_WORKSPACE" +# echo "Runner Temp Dir: $RUNNER_TEMP" +# echo "Runner Tool Cache: $RUNNER_TOOL_CACHE" +# - name: Install Primus +# run: | +# pip install -r requirements-jax.txt +# - name: Set UT_LOG_PATH +# run: | +# ts="$(date +%Y%m%d-%H%M%S)" +# commit_id="${GITHUB_SHA::7}" +# if [[ "${{ github.event_name }}" == "pull_request" ]]; then +# 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 }}" == "release" ]]; then +# TAG_NAME="${{ github.ref }}" +# TAG="${TAG_NAME#refs/tags/}" +# echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/${TAG}-${ts}-${commit_id}" >> $GITHUB_ENV +# else +# echo "UT_LOG_PATH=${PRIMUS_WORKDIR}/ut_out/others-${ts}-${commit_id}" >> $GITHUB_ENV +# fi +# - name: Run Shell Tests +# run: | +# echo "Running Primus CLI shell tests..." +# bash ./tests/runner/run_all_tests.sh +# - name: Setup MaxText E2E coverage +# if: always() +# continue-on-error: true +# run: | +# # Instrument the MaxText (JAX) training subprocesses so their coverage +# # of primus/ (notably backends/maxtext) is recorded. Same mechanism as +# # the torch job: a .pth runs coverage.process_startup() in every child. +# pip install coverage >/dev/null 2>&1 || true +# 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" +# - name: Decide E2E scope (jax) +# run: | +# # Fail-safe: default to running MaxText E2E; on PRs skip it when no +# # relevant (maxtext / shared) changes are detected. +# X=1 +# 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)" +# echo "Selected jax E2E scope: ${e2e:-}" +# if [[ "${e2e}" != "all" ]]; then +# echo "${e2e}" | grep -qw maxtext || X=0 +# fi +# fi +# fi +# echo "RUN_MAXTEXT_E2E=${X}" >> "$GITHUB_ENV" +# - name: Run Unit Tests +# env: +# HF_TOKEN: ${{secrets.HF_TOKEN}} +# run: | +# if [[ "${RUN_MAXTEXT_E2E:-1}" != "1" ]]; then echo "Skipping MaxText 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 }}" +# mkdir -p "${GITHUB_WORKSPACE}/test-reports" +# # 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" +# # 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 +# - name: Upload test reports +# if: always() +# uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 +# with: +# name: jax-test-reports +# path: test-reports/ +# if-no-files-found: warn +# retention-days: 14 +# - name: Write test summary +# if: always() +# run: | +# python tools/ci/junit_summary.py --title jax test-reports/*.xml >> "$GITHUB_STEP_SUMMARY" || true +# - name: Write runtime summary +# if: always() +# run: | +# python tools/ci/runtime_summary.py --title jax "$RUNNER_TEMP/runtime.tsv" >> "$GITHUB_STEP_SUMMARY" || true +# - name: Build jax coverage json (MaxText E2E) +# if: always() +# continue-on-error: true +# run: | +# unset COVERAGE_PROCESS_START +# # The JSON is consumed by the coverage-summary job (no per-job summary). +# python -m coverage combine 2>/dev/null || true +# python -m coverage json -o coverage_maxtext_e2e.json 2>/dev/null || true +# - name: Upload jax coverage json +# if: always() +# continue-on-error: true +# uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 +# with: +# name: coverage-jax +# path: coverage_maxtext_e2e.json +# if-no-files-found: ignore +# - name: Clean +# if: always() +# run: | +# rm -rf ${PRIMUS_WORKDIR}/Primus-Turbo +# rm -rf ${PRIMUS_WORKDIR}/Primus +# # Remove E2E coverage artifacts + the injection .pth (persistent runner). +# rm -f "${GITHUB_WORKSPACE}/.coverage_e2e"* "${GITHUB_WORKSPACE}/.coveragerc_e2e" coverage_maxtext_e2e.json || true +# SP=$(python -c "import site; print(site.getsitepackages()[0])" 2>/dev/null) && rm -f "$SP/primus_e2e_coverage.pth" || true + +# coverage-summary: +# # Combine the torch and jax coverage JSON artifacts into one unit-vs-E2E +# # table. Uses module-level JSON (not coverage data), so no cross-runner path +# # mapping is needed: torch (megatron/torchtitan) and jax (maxtext) E2E cover +# # near-disjoint modules and are merged per module by max covered lines. +# needs: [run-unittest-torch, run-unittest-jax] +# if: always() +# runs-on: ubuntu-latest +# steps: +# - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 +# - name: Download coverage artifacts +# uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 +# continue-on-error: true +# with: +# path: cov-artifacts +# - name: Render combined coverage summary +# continue-on-error: true +# run: | +# unit=$(find cov-artifacts -name coverage_unit.json | head -1) +# torch=$(find cov-artifacts -name coverage_combined.json | head -1) +# jax=$(find cov-artifacts -name coverage_maxtext_e2e.json | head -1) +# if [ -n "$unit" ]; then +# python tools/ci/coverage_summary.py "$unit" "Unit vs Unit+E2E (torch + jax)" $torch $jax >> "$GITHUB_STEP_SUMMARY" +# else +# echo "No unit coverage artifact found; skipping combined summary." >> "$GITHUB_STEP_SUMMARY" +# fi diff --git a/.github/workflows/deploy-projection.yml b/.github/workflows/deploy-projection.yml index 5561a4168..aee238d51 100644 --- a/.github/workflows/deploy-projection.yml +++ b/.github/workflows/deploy-projection.yml @@ -12,7 +12,7 @@ on: workflow_dispatch: push: branches: - - dev/wenx/deepseek-v4 + - dev/tas/deepseek-v4 paths: - "deepseek-v4/projection/site/**" - ".github/workflows/deploy-projection.yml" diff --git a/.github/workflows/docker/Dockerfile b/.github/workflows/docker/Dockerfile index d339060cc..08a825f3c 100644 --- a/.github/workflows/docker/Dockerfile +++ b/.github/workflows/docker/Dockerfile @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=docker.io/rocm/primus:v26.2 +ARG BASE_IMAGE=docker.io/rocm/primus:v26.3 FROM ${BASE_IMAGE} ARG PRIMUS_TURBO_COMMIT @@ -23,7 +23,7 @@ RUN rm -rf /var/lib/apt/lists/* ENV ROCSHMEM_HOME=/opt/rocshmem ENV UCX_HOME=/opt/ucx # ENV MPI_HOME=/opt/ompi -# Use the system OpenMPI prefix from the v26.2 base image. +# 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} @@ -55,9 +55,10 @@ RUN rm -rf /opt/rocSHMEM RUN pip3 uninstall aiter amd-aiter -y || true && \ cd /opt && \ - git clone --recursive https://github.com/ROCm/aiter.git && \ + git clone https://github.com/ROCm/aiter.git && \ cd aiter && \ - git checkout ${PRIMUS_TURBO_AITER_COMMIT} && \ + git checkout -f ${PRIMUS_TURBO_AITER_COMMIT} && \ + git submodule sync && \ git submodule update --init --recursive && \ PREBUILD_KERNELS=3 pip install --no-cache-dir --use-pep517 . @@ -71,6 +72,18 @@ RUN cd /opt && \ RUN rm -rf /opt/Primus-Turbo +# --------------------------------------------------------------------------- +# Install Triton +# --------------------------------------------------------------------------- +RUN cd /opt && \ + git clone https://github.com/triton-lang/triton.git && \ + cd triton && \ + git checkout 09500db9 && \ + pip3 install -r python/requirements.txt && \ + MAX_JOBS=96 pip3 install --no-build-isolation --force-reinstall --no-deps . + +RUN rm -rf /opt/triton + # --------------------------------------------------------------------------- # Install UCCL-EP (skip for JAX framework) # --------------------------------------------------------------------------- @@ -86,6 +99,24 @@ 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. +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 && \ + 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 --no-cache-dir ./shared/origami/python && \ + rm -rf /tmp/rocm-libraries; \ + fi + # Set the default working directory WORKDIR /opt diff --git a/.github/workflows/docker/Dockerfile.ainic b/.github/workflows/docker/Dockerfile.ainic index f377bd5d9..ca103df2d 100644 --- a/.github/workflows/docker/Dockerfile.ainic +++ b/.github/workflows/docker/Dockerfile.ainic @@ -19,9 +19,9 @@ RUN apt-get update && \ # --------------------------------------------------------------------------- # Enviroment variables # --------------------------------------------------------------------------- -ENV WORKDIR=/opt +ENV WORKDIR=/workspace ENV ROCM_PATH=/opt/rocm -ENV MPI_PATH=/opt/ompi +ENV MPI_PATH=/usr/lib/x86_64-linux-gnu/openmpi # =============================== Build AINIC Driver =============================== # WARNING: Please ensure the following environment variables are correctly set: @@ -30,16 +30,24 @@ ENV MPI_PATH=/opt/ompi # 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} +COPY ${AINIC_BUNDLE_PATH}/ainic_bundle_1.117.5-a-77.tar.gz ${WORKDIR} RUN cd ${WORKDIR} && \ + rm -rf rccl amd-anp && \ 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 ainic_bundle_1.117.5-a-77.tar.gz && \ + cd ainic_bundle_1.117.5-a-77 && \ 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 + printf '%s\n' \ + 'Package: libionic1 libionic-dev' \ + 'Pin: version 54.0-187-1' \ + 'Pin-Priority: 1001' \ + '' \ + 'Package: perftest' \ + 'Pin: version 1:25.04.0.0.84-128-1' \ + 'Pin-Priority: 1001' \ + > /etc/apt/preferences.d/ainic-a77 && \ + ./install.sh --domain=user -y 2>&1 | tee log_install.txt # --------------------------------------------------------------------------- # Build rccl diff --git a/.gitignore b/.gitignore index 11a4d0c4d..6e93e90d2 100644 --- a/.gitignore +++ b/.gitignore @@ -17,3 +17,8 @@ pp_simulation_result # GitHub Pages) despite the generic `data` ignore above. !deepseek-v4/projection/site/data/ !deepseek-v4/projection/site/data/*.json + +*.log +*.nohup +*.zip +.triton_cache_shared/ diff --git a/check_hc_expand.py b/check_hc_expand.py new file mode 100644 index 000000000..481303ceb --- /dev/null +++ b/check_hc_expand.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +"""Correctness + microbench for the fused HC expand Triton kernel vs eager.""" +import importlib.util +import os +import time + +import torch + +os.environ.setdefault("PRIMUS_HC_EXPAND_TRITON", "1") +# Load the kernel module directly by path to avoid the heavy primus/__init__ chain. +_MOD = os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/hc_expand.py", +) +_spec = importlib.util.spec_from_file_location("hc_expand", _MOD) +_hc = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_hc) +hc_expand_triton = _hc.hc_expand_triton + + +def eager_expand(x, out, post, comb): + write = post.unsqueeze(-1) * out.unsqueeze(-2) + mix = torch.matmul(comb, x) + return write + mix + + +def run(leading, K, D, dtype): + dev = "cuda" + torch.manual_seed(0) + mk = lambda *s: torch.randn(*s, device=dev, dtype=dtype) # noqa: E731 + x = mk(*leading, K, D) + out = mk(*leading, D) + post = mk(*leading, K) + comb = torch.softmax(mk(*leading, K, K).float(), dim=-1).to(dtype) + + # fp32 reference (ground truth), plus eager and triton at the test dtype. + x32, o32, p32, c32 = (t.float().clone().requires_grad_(True) for t in (x, out, post, comb)) + xa, oa, pa, ca = (t.clone().requires_grad_(True) for t in (x, out, post, comb)) + xb, ob, pb, cb = (t.clone().requires_grad_(True) for t in (x, out, post, comb)) + + ref = eager_expand(x32, o32, p32, c32) + eag = eager_expand(xa, oa, pa, ca) + tri = hc_expand_triton(xb, ob, pb, cb) + + g32 = torch.randn_like(ref) + g = g32.to(dtype) + ref.backward(g32) + eag.backward(g) + tri.backward(g) + + def maxerr(a, b): + return (a.float() - b.float()).abs().max().item() + + tag = f"{tuple(leading)} K={K} D={D} {str(dtype).split('.')[-1]}" + ok = True + for name, r, e, t in [ + ("fwd", ref, eag, tri), + ("dx", x32.grad, xa.grad, xb.grad), + ("dout", o32.grad, oa.grad, ob.grad), + ("dpost", p32.grad, pa.grad, pb.grad), + ("dcomb", c32.grad, ca.grad, cb.grad), + ]: + e_err = maxerr(r, e) # eager-vs-fp32 (the dtype's intrinsic noise floor) + t_err = maxerr(r, t) # triton-vs-fp32 + # Pass if triton is no worse than eager, with an fp32 reduction-order + # slack relative to the tensor magnitude (length-D sums reorder). + slack = 1e-4 + 5e-5 * r.abs().max().item() + passed = t_err <= max(1.5 * e_err, slack) + ok = ok and passed + print(f" {name:6s} triton_err={t_err:.3e} eager_err={e_err:.3e} {'OK' if passed else 'FAIL'}") + print(f"[{tag}] {'PASS' if ok else 'FAIL'}") + return ok + + +def bench(leading, K, D, dtype, iters=50): + dev = "cuda" + x = torch.randn(*leading, K, D, device=dev, dtype=dtype, requires_grad=True) + out = torch.randn(*leading, D, device=dev, dtype=dtype, requires_grad=True) + post = torch.randn(*leading, K, device=dev, dtype=dtype, requires_grad=True) + comb = ( + torch.softmax(torch.randn(*leading, K, K, device=dev).float(), dim=-1).to(dtype).requires_grad_(True) + ) + g = torch.randn(*leading, K, D, device=dev, dtype=dtype) + + def step(fn): + for t in (x, out, post, comb): + t.grad = None + y = fn(x, out, post, comb) + y.backward(g) + + for fn, name in [(eager_expand, "eager"), (hc_expand_triton, "triton")]: + step(fn) + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(iters): + step(fn) + torch.cuda.synchronize() + ms = (time.perf_counter() - t0) / iters * 1e3 + print(f" {name:7s} fwd+bwd {ms:8.3f} ms/iter") + + +if __name__ == "__main__": + print("=== correctness ===") + all_ok = True + all_ok &= run((4096,), 4, 7168, torch.bfloat16) + all_ok &= run((4096,), 4, 7168, torch.float32) + all_ok &= run((2, 17), 4, 320, torch.bfloat16) + all_ok &= run((3,), 2, 128, torch.float32) + all_ok &= run((5,), 8, 64, torch.bfloat16) + print("\n=== bench (production shape B*S=4096, K=4, D=7168, bf16) ===") + bench((4096,), 4, 7168, torch.bfloat16) + print("\nALL", "PASS" if all_ok else "FAIL") diff --git a/deepseek-v4/benchmark/bench_v4_attention.py b/deepseek-v4/benchmark/bench_v4_attention.py new file mode 100644 index 000000000..9a5b5e516 --- /dev/null +++ b/deepseek-v4/benchmark/bench_v4_attention.py @@ -0,0 +1,619 @@ +############################################################################### +# 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 deepseek-v4/benchmark/bench_v4_attention.py + python 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 +try: + _flydsl_turbo_dir = os.environ.get( + "PRIMUS_V4_FLYDSL_TURBO_DIR", + os.path.join( + os.path.dirname(__file__), + "..", + "..", + "agent", + "workspace", + "flydsl_turbo_dsv4_20260707", + ), + ) + _flydsl_turbo_dir = os.path.abspath(_flydsl_turbo_dir) + if _flydsl_turbo_dir not in sys.path: + sys.path.insert(0, _flydsl_turbo_dir) + from flydsl_turbo_adapter import ( + sparse_mla_bwd_v4_flydsl_turbo, + sparse_mla_fwd_v4_flydsl_turbo, + ) + + _SPARSE_MLA_BACKENDS["turbo_flydsl"] = ( + sparse_mla_fwd_v4_flydsl_turbo, + sparse_mla_bwd_v4_flydsl_turbo, + ) +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", + "_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/deepseek-v4/benchmark/bench_v4_attention_results.md b/deepseek-v4/benchmark/bench_v4_attention_results.md new file mode 100644 index 000000000..54a99be15 --- /dev/null +++ b/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 (`smci355-ccs-aus-n04-25`) +- **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= python deepseek-v4/benchmark/bench_v4_attention.py \ + --variant both --cr all --warmup 10 --iters 30 +``` diff --git a/deepseek-v4/develop/plan-4/README.md b/deepseek-v4/develop/plan-4/README.md index 9205beb2e..ad36ef3f5 100644 --- a/deepseek-v4/develop/plan-4/README.md +++ b/deepseek-v4/develop/plan-4/README.md @@ -96,5 +96,3 @@ the release gate). | **P25** | `v4_attention` Triton kernel (forward + backward) for `compress_ratio ∈ {0, 128}` + `use_v4_triton_attention` switch | "My understanding: develop two attention versions — one normal attention used for compress=0 and 128 — implemented in Triton." | | **P26** | `v4_csa_attention` Triton kernel (forward + backward) for `compress_ratio == 4` + `use_v4_triton_csa_attention` switch | "Plus a CSA version — use the two switches `use_v4_triton_attention` and `use_v4_triton_csa_attention` to control whether the new Triton versions are used." | | **P27** | V4 attention dispatch wiring + `run_deepseek_v4.sh` smoke (TP=1 PP=1 EP=8) | "The PyTorch small-op version is currently in tree — keep it. You can add new Triton versions and use `use_v4_triton_attention` / `use_v4_triton_csa_attention` switches to control which is used." | - - diff --git a/deepseek-v4/develop/progress/p19/run_profile_ep8.sh b/deepseek-v4/develop/progress/p19/run_profile_ep8.sh index 95980d158..d3db1a016 100755 --- a/deepseek-v4/develop/progress/p19/run_profile_ep8.sh +++ b/deepseek-v4/develop/progress/p19/run_profile_ep8.sh @@ -44,9 +44,11 @@ export FP8=null export FP8_RECIPE=null export EXP=examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml -export BACKEND_PATH="$(pwd)/third_party/Megatron-LM" +BACKEND_PATH="$(pwd)/third_party/Megatron-LM" +export BACKEND_PATH export PRIMUS_TEAM=amd -export PRIMUS_USER=tas-mi355x-$(date +%Y%m%d) +PRIMUS_USER=tas-mi355x-$(date +%Y%m%d) +export PRIMUS_USER export PRIMUS_EXP_NAME=p19_profile_pp1_ep8 mkdir -p "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" diff --git a/deepseek-v4/develop/progress/p19/run_profile_pp2_ep4.sh b/deepseek-v4/develop/progress/p19/run_profile_pp2_ep4.sh index 535fa078b..8360b0ce4 100755 --- a/deepseek-v4/develop/progress/p19/run_profile_pp2_ep4.sh +++ b/deepseek-v4/develop/progress/p19/run_profile_pp2_ep4.sh @@ -44,9 +44,11 @@ export FP8=null export FP8_RECIPE=null export EXP=examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml -export BACKEND_PATH="$(pwd)/third_party/Megatron-LM" +BACKEND_PATH="$(pwd)/third_party/Megatron-LM" +export BACKEND_PATH export PRIMUS_TEAM=amd -export PRIMUS_USER=tas-mi355x-$(date +%Y%m%d) +PRIMUS_USER=tas-mi355x-$(date +%Y%m%d) +export PRIMUS_USER export PRIMUS_EXP_NAME=p19_profile_pp2_ep4 mkdir -p "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" diff --git a/deepseek-v4/develop/progress/p19/run_smokeC_v2.sh b/deepseek-v4/develop/progress/p19/run_smokeC_v2.sh index 677ee92d3..0251b7049 100755 --- a/deepseek-v4/develop/progress/p19/run_smokeC_v2.sh +++ b/deepseek-v4/develop/progress/p19/run_smokeC_v2.sh @@ -43,9 +43,11 @@ export FP8=null export FP8_RECIPE=null export EXP=examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml -export BACKEND_PATH="$(pwd)/third_party/Megatron-LM" +BACKEND_PATH="$(pwd)/third_party/Megatron-LM" +export BACKEND_PATH export PRIMUS_TEAM=amd -export PRIMUS_USER=tas-mi355x-$(date +%Y%m%d) +PRIMUS_USER=tas-mi355x-$(date +%Y%m%d) +export PRIMUS_USER export PRIMUS_EXP_NAME=p19_smokeC_pp4_ep2_v2 mkdir -p "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" diff --git a/deepseek-v4/develop/progress/p19/run_smokeD_v2.sh b/deepseek-v4/develop/progress/p19/run_smokeD_v2.sh index 42b6fc11e..0d8a760fd 100755 --- a/deepseek-v4/develop/progress/p19/run_smokeD_v2.sh +++ b/deepseek-v4/develop/progress/p19/run_smokeD_v2.sh @@ -44,9 +44,11 @@ export FP8=null export FP8_RECIPE=null export EXP=examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml -export BACKEND_PATH="$(pwd)/third_party/Megatron-LM" +BACKEND_PATH="$(pwd)/third_party/Megatron-LM" +export BACKEND_PATH export PRIMUS_TEAM=amd -export PRIMUS_USER=tas-mi355x-$(date +%Y%m%d) +PRIMUS_USER=tas-mi355x-$(date +%Y%m%d) +export PRIMUS_USER export PRIMUS_EXP_NAME=p19_smokeD_pp2_ep4_vpp2_v2 mkdir -p "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" diff --git a/deepseek-v4/develop/progress/p20/smoke_ep8_pp1_final.log b/deepseek-v4/develop/progress/p20/smoke_ep8_pp1_final.log index 3bcf7e068..308e7f160 100644 --- a/deepseek-v4/develop/progress/p20/smoke_ep8_pp1_final.log +++ b/deepseek-v4/develop/progress/p20/smoke_ep8_pp1_final.log @@ -120,7 +120,7 @@ uv.lock' ']' [2026-05-07 13:07:21] [NODE-?(mi355-gpu-12)] [INFO] ----------------------------------------------- [2026-05-07 13:07:21] [NODE-?(mi355-gpu-12)] [INFO] primus-cli-direct.sh [2026-05-07 13:07:21] [NODE-?(mi355-gpu-12)] [INFO] ----------------------------------------------- -[2026-05-07 13:07:21] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRE_PARSE_ARGS (runner options): +[2026-05-07 13:07:21] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRE_PARSE_ARGS (runner options): [2026-05-07 13:07:21] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRIMUS_ARGS (python module args): train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 1 --expert_model_parallel_size 8 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --use_turbo_attention False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 [2026-05-07 13:07:21] [NODE-?(mi355-gpu-12)] [INFO] [direct] Combined args for second pass: -- train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 1 --expert_model_parallel_size 8 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --use_turbo_attention False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 [2026-05-07 13:07:21] [NODE-?(mi355-gpu-12)] [INFO] [direct] Loaded config: /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/runner/.primus.yaml @@ -325,11 +325,11 @@ WARNING: AMD GPU device(s) is/are in a low-power state. Check power control/runt Executing: torchrun --nproc_per_node 8 --nnodes 1 --node_rank 0 --master_addr localhost --master_port 1234 --local-ranks-filter 0 primus/cli/main.py train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 1 --expert_model_parallel_size 8 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --use_turbo_attention False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM 2>&1 | tee logs/log_20260507_130721.txt ========================================== - + ========================================== -W0507 13:07:24.408000 62629 torch/distributed/run.py:852] +W0507 13:07:24.408000 62629 torch/distributed/run.py:852] W0507 13:07:24.408000 62629 torch/distributed/run.py:852] ***************************************** -W0507 13:07:24.408000 62629 torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. +W0507 13:07:24.408000 62629 torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. W0507 13:07:24.408000 62629 torch/distributed/run.py:852] ***************************************** [Primus:Runtime] Applying CLI overrides: {'num_layers': 8, 'train_iters': 10, 'lr_warmup_iters': 0, 'lr_decay_iters': 10, 'micro_batch_size': 1, 'global_batch_size': 16, 'seq_length': 128, 'max_position_embeddings': 128, 'rope_type': 'rope', 'tensor_model_parallel_size': 1, 'pipeline_model_parallel_size': 1, 'expert_model_parallel_size': 8, 'num_experts': 8, 'moe_router_topk': 2, 'moe_router_enable_expert_bias': False, 'moe_ffn_hidden_size': 512, 'index_topk': 8, 'v4_grouped_experts_support_clamped_swiglu': True, 'compress_ratios': '[0,0,4,4,4,4,4,0]', 'mtp_num_layers': 0, 'mock_data': True, 'use_turbo_attention': False, 'use_turbo_grouped_mlp': False, 'moe_use_legacy_grouped_gemm': False, 'fp8': 'null', 'fp8_recipe': 'null', 'recompute_num_layers': 0, 'recompute_granularity': 'full', 'recompute_method': 'block', 'overlap_grad_reduce': False, 'overlap_param_gather': False, 'disable_last_saving': True, 'disable_wandb': True, 'disable_tensorboard': True, 'profile': False, 'use_pytorch_profiler': False, 'profile_step_end': 7, 'profile_step_start': 6} [Primus:Env] HF_HOME already set: /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/data/huggingface diff --git a/deepseek-v4/develop/progress/p20/smoke_pp2_ep4_final.log b/deepseek-v4/develop/progress/p20/smoke_pp2_ep4_final.log index fe1ff0169..626493372 100644 --- a/deepseek-v4/develop/progress/p20/smoke_pp2_ep4_final.log +++ b/deepseek-v4/develop/progress/p20/smoke_pp2_ep4_final.log @@ -120,7 +120,7 @@ uv.lock' ']' [2026-05-07 13:09:21] [NODE-?(mi355-gpu-12)] [INFO] ----------------------------------------------- [2026-05-07 13:09:21] [NODE-?(mi355-gpu-12)] [INFO] primus-cli-direct.sh [2026-05-07 13:09:21] [NODE-?(mi355-gpu-12)] [INFO] ----------------------------------------------- -[2026-05-07 13:09:21] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRE_PARSE_ARGS (runner options): +[2026-05-07 13:09:21] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRE_PARSE_ARGS (runner options): [2026-05-07 13:09:21] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRIMUS_ARGS (python module args): train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 2 --expert_model_parallel_size 4 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --use_turbo_attention False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 [2026-05-07 13:09:21] [NODE-?(mi355-gpu-12)] [INFO] [direct] Combined args for second pass: -- train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 2 --expert_model_parallel_size 4 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --use_turbo_attention False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 [2026-05-07 13:09:21] [NODE-?(mi355-gpu-12)] [INFO] [direct] Loaded config: /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/runner/.primus.yaml @@ -325,11 +325,11 @@ WARNING: AMD GPU device(s) is/are in a low-power state. Check power control/runt Executing: torchrun --nproc_per_node 8 --nnodes 1 --node_rank 0 --master_addr localhost --master_port 1234 --local-ranks-filter 0 primus/cli/main.py train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 2 --expert_model_parallel_size 4 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --use_turbo_attention False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM 2>&1 | tee logs/log_20260507_130921.txt ========================================== - + ========================================== -W0507 13:09:24.676000 65128 torch/distributed/run.py:852] +W0507 13:09:24.676000 65128 torch/distributed/run.py:852] W0507 13:09:24.676000 65128 torch/distributed/run.py:852] ***************************************** -W0507 13:09:24.676000 65128 torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. +W0507 13:09:24.676000 65128 torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. W0507 13:09:24.676000 65128 torch/distributed/run.py:852] ***************************************** [Primus:Runtime] Applying CLI overrides: {'num_layers': 8, 'train_iters': 10, 'lr_warmup_iters': 0, 'lr_decay_iters': 10, 'micro_batch_size': 1, 'global_batch_size': 16, 'seq_length': 128, 'max_position_embeddings': 128, 'rope_type': 'rope', 'tensor_model_parallel_size': 1, 'pipeline_model_parallel_size': 2, 'expert_model_parallel_size': 4, 'num_experts': 8, 'moe_router_topk': 2, 'moe_router_enable_expert_bias': False, 'moe_ffn_hidden_size': 512, 'index_topk': 8, 'v4_grouped_experts_support_clamped_swiglu': True, 'compress_ratios': '[0,0,4,4,4,4,4,0]', 'mtp_num_layers': 0, 'mock_data': True, 'use_turbo_attention': False, 'use_turbo_grouped_mlp': False, 'moe_use_legacy_grouped_gemm': False, 'fp8': 'null', 'fp8_recipe': 'null', 'recompute_num_layers': 0, 'recompute_granularity': 'full', 'recompute_method': 'block', 'overlap_grad_reduce': False, 'overlap_param_gather': False, 'disable_last_saving': True, 'disable_wandb': True, 'disable_tensorboard': True, 'profile': False, 'use_pytorch_profiler': False, 'profile_step_end': 7, 'profile_step_start': 6} [Primus:Env] HF_HOME already set: /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/data/huggingface diff --git a/deepseek-v4/develop/progress/p21/smoke_ep8_pp1.log b/deepseek-v4/develop/progress/p21/smoke_ep8_pp1.log index 0320a1186..4408d8d75 100644 --- a/deepseek-v4/develop/progress/p21/smoke_ep8_pp1.log +++ b/deepseek-v4/develop/progress/p21/smoke_ep8_pp1.log @@ -120,7 +120,7 @@ uv.lock' ']' [2026-05-07 13:45:03] [NODE-?(mi355-gpu-12)] [INFO] ----------------------------------------------- [2026-05-07 13:45:03] [NODE-?(mi355-gpu-12)] [INFO] primus-cli-direct.sh [2026-05-07 13:45:03] [NODE-?(mi355-gpu-12)] [INFO] ----------------------------------------------- -[2026-05-07 13:45:03] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRE_PARSE_ARGS (runner options): +[2026-05-07 13:45:03] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRE_PARSE_ARGS (runner options): [2026-05-07 13:45:03] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRIMUS_ARGS (python module args): train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 1 --expert_model_parallel_size 8 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --use_turbo_attention False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 [2026-05-07 13:45:03] [NODE-?(mi355-gpu-12)] [INFO] [direct] Combined args for second pass: -- train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 1 --expert_model_parallel_size 8 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --use_turbo_attention False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 [2026-05-07 13:45:03] [NODE-?(mi355-gpu-12)] [INFO] [direct] Loaded config: /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/runner/.primus.yaml @@ -325,11 +325,11 @@ WARNING: AMD GPU device(s) is/are in a low-power state. Check power control/runt Executing: torchrun --nproc_per_node 8 --nnodes 1 --node_rank 0 --master_addr localhost --master_port 1234 --local-ranks-filter 0 primus/cli/main.py train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 1 --expert_model_parallel_size 8 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --use_turbo_attention False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM 2>&1 | tee logs/log_20260507_134503.txt ========================================== - + ========================================== -W0507 13:45:06.283000 69458 torch/distributed/run.py:852] +W0507 13:45:06.283000 69458 torch/distributed/run.py:852] W0507 13:45:06.283000 69458 torch/distributed/run.py:852] ***************************************** -W0507 13:45:06.283000 69458 torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. +W0507 13:45:06.283000 69458 torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. W0507 13:45:06.283000 69458 torch/distributed/run.py:852] ***************************************** [Primus:Runtime] Applying CLI overrides: {'num_layers': 8, 'train_iters': 10, 'lr_warmup_iters': 0, 'lr_decay_iters': 10, 'micro_batch_size': 1, 'global_batch_size': 16, 'seq_length': 128, 'max_position_embeddings': 128, 'rope_type': 'rope', 'tensor_model_parallel_size': 1, 'pipeline_model_parallel_size': 1, 'expert_model_parallel_size': 8, 'num_experts': 8, 'moe_router_topk': 2, 'moe_router_enable_expert_bias': False, 'moe_ffn_hidden_size': 512, 'index_topk': 8, 'v4_grouped_experts_support_clamped_swiglu': True, 'compress_ratios': '[0,0,4,4,4,4,4,0]', 'mtp_num_layers': 0, 'mock_data': True, 'use_turbo_attention': False, 'use_turbo_grouped_mlp': False, 'moe_use_legacy_grouped_gemm': False, 'fp8': 'null', 'fp8_recipe': 'null', 'recompute_num_layers': 0, 'recompute_granularity': 'full', 'recompute_method': 'block', 'overlap_grad_reduce': False, 'overlap_param_gather': False, 'disable_last_saving': True, 'disable_wandb': True, 'disable_tensorboard': True, 'profile': False, 'use_pytorch_profiler': False, 'profile_step_end': 7, 'profile_step_start': 6} [Primus:Env] HF_HOME already set: /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/data/huggingface diff --git a/deepseek-v4/develop/progress/p21/smoke_pp2_ep4.log b/deepseek-v4/develop/progress/p21/smoke_pp2_ep4.log index 773e239eb..1065c690d 100644 --- a/deepseek-v4/develop/progress/p21/smoke_pp2_ep4.log +++ b/deepseek-v4/develop/progress/p21/smoke_pp2_ep4.log @@ -120,7 +120,7 @@ uv.lock' ']' [2026-05-07 13:46:49] [NODE-?(mi355-gpu-12)] [INFO] ----------------------------------------------- [2026-05-07 13:46:49] [NODE-?(mi355-gpu-12)] [INFO] primus-cli-direct.sh [2026-05-07 13:46:49] [NODE-?(mi355-gpu-12)] [INFO] ----------------------------------------------- -[2026-05-07 13:46:49] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRE_PARSE_ARGS (runner options): +[2026-05-07 13:46:49] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRE_PARSE_ARGS (runner options): [2026-05-07 13:46:49] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRIMUS_ARGS (python module args): train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 2 --expert_model_parallel_size 4 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --use_turbo_attention False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 [2026-05-07 13:46:49] [NODE-?(mi355-gpu-12)] [INFO] [direct] Combined args for second pass: -- train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 2 --expert_model_parallel_size 4 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --use_turbo_attention False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 [2026-05-07 13:46:49] [NODE-?(mi355-gpu-12)] [INFO] [direct] Loaded config: /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/runner/.primus.yaml @@ -325,11 +325,11 @@ WARNING: AMD GPU device(s) is/are in a low-power state. Check power control/runt Executing: torchrun --nproc_per_node 8 --nnodes 1 --node_rank 0 --master_addr localhost --master_port 1234 --local-ranks-filter 0 primus/cli/main.py train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 2 --expert_model_parallel_size 4 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --use_turbo_attention False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM 2>&1 | tee logs/log_20260507_134649.txt ========================================== - + ========================================== -W0507 13:46:52.343000 71957 torch/distributed/run.py:852] +W0507 13:46:52.343000 71957 torch/distributed/run.py:852] W0507 13:46:52.343000 71957 torch/distributed/run.py:852] ***************************************** -W0507 13:46:52.343000 71957 torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. +W0507 13:46:52.343000 71957 torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. W0507 13:46:52.343000 71957 torch/distributed/run.py:852] ***************************************** [Primus:Runtime] Applying CLI overrides: {'num_layers': 8, 'train_iters': 10, 'lr_warmup_iters': 0, 'lr_decay_iters': 10, 'micro_batch_size': 1, 'global_batch_size': 16, 'seq_length': 128, 'max_position_embeddings': 128, 'rope_type': 'rope', 'tensor_model_parallel_size': 1, 'pipeline_model_parallel_size': 2, 'expert_model_parallel_size': 4, 'num_experts': 8, 'moe_router_topk': 2, 'moe_router_enable_expert_bias': False, 'moe_ffn_hidden_size': 512, 'index_topk': 8, 'v4_grouped_experts_support_clamped_swiglu': True, 'compress_ratios': '[0,0,4,4,4,4,4,0]', 'mtp_num_layers': 0, 'mock_data': True, 'use_turbo_attention': False, 'use_turbo_grouped_mlp': False, 'moe_use_legacy_grouped_gemm': False, 'fp8': 'null', 'fp8_recipe': 'null', 'recompute_num_layers': 0, 'recompute_granularity': 'full', 'recompute_method': 'block', 'overlap_grad_reduce': False, 'overlap_param_gather': False, 'disable_last_saving': True, 'disable_wandb': True, 'disable_tensorboard': True, 'profile': False, 'use_pytorch_profiler': False, 'profile_step_end': 7, 'profile_step_start': 6} [Primus:Env] HF_HOME already set: /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/data/huggingface diff --git a/deepseek-v4/develop/progress/p22/smoke_eager_ep8_pp1.log b/deepseek-v4/develop/progress/p22/smoke_eager_ep8_pp1.log index 5c18189fc..10a2e407f 100644 --- a/deepseek-v4/develop/progress/p22/smoke_eager_ep8_pp1.log +++ b/deepseek-v4/develop/progress/p22/smoke_eager_ep8_pp1.log @@ -126,7 +126,7 @@ uv.lock' ']' [2026-05-08 02:18:42] [NODE-?(mi355-gpu-12)] [INFO] ----------------------------------------------- [2026-05-08 02:18:42] [NODE-?(mi355-gpu-12)] [INFO] primus-cli-direct.sh [2026-05-08 02:18:42] [NODE-?(mi355-gpu-12)] [INFO] ----------------------------------------------- -[2026-05-08 02:18:42] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRE_PARSE_ARGS (runner options): +[2026-05-08 02:18:42] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRE_PARSE_ARGS (runner options): [2026-05-08 02:18:42] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRIMUS_ARGS (python module args): train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 1 --expert_model_parallel_size 8 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --enable_primus_turbo False --use_turbo_attention False --use_turbo_deepep False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 [2026-05-08 02:18:42] [NODE-?(mi355-gpu-12)] [INFO] [direct] Combined args for second pass: -- train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 1 --expert_model_parallel_size 8 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --enable_primus_turbo False --use_turbo_attention False --use_turbo_deepep False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 [2026-05-08 02:18:42] [NODE-?(mi355-gpu-12)] [INFO] [direct] Loaded config: /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/runner/.primus.yaml @@ -331,11 +331,11 @@ WARNING: AMD GPU device(s) is/are in a low-power state. Check power control/runt Executing: torchrun --nproc_per_node 8 --nnodes 1 --node_rank 0 --master_addr localhost --master_port 1234 --local-ranks-filter 0 primus/cli/main.py train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 1 --expert_model_parallel_size 8 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --enable_primus_turbo False --use_turbo_attention False --use_turbo_deepep False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM 2>&1 | tee logs/log_20260508_021842.txt ========================================== - + ========================================== -W0508 02:18:45.387000 94465 torch/distributed/run.py:852] +W0508 02:18:45.387000 94465 torch/distributed/run.py:852] W0508 02:18:45.387000 94465 torch/distributed/run.py:852] ***************************************** -W0508 02:18:45.387000 94465 torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. +W0508 02:18:45.387000 94465 torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. W0508 02:18:45.387000 94465 torch/distributed/run.py:852] ***************************************** [Primus:Runtime] Applying CLI overrides: {'num_layers': 8, 'train_iters': 10, 'lr_warmup_iters': 0, 'lr_decay_iters': 10, 'micro_batch_size': 1, 'global_batch_size': 16, 'seq_length': 128, 'max_position_embeddings': 128, 'rope_type': 'rope', 'tensor_model_parallel_size': 1, 'pipeline_model_parallel_size': 1, 'expert_model_parallel_size': 8, 'num_experts': 8, 'moe_router_topk': 2, 'moe_router_enable_expert_bias': False, 'moe_ffn_hidden_size': 512, 'index_topk': 8, 'v4_grouped_experts_support_clamped_swiglu': True, 'compress_ratios': '[0,0,4,4,4,4,4,0]', 'mtp_num_layers': 0, 'mock_data': True, 'enable_primus_turbo': False, 'use_turbo_attention': False, 'use_turbo_deepep': False, 'use_turbo_grouped_mlp': False, 'moe_use_legacy_grouped_gemm': False, 'fp8': 'null', 'fp8_recipe': 'null', 'recompute_num_layers': 0, 'recompute_granularity': 'full', 'recompute_method': 'block', 'overlap_grad_reduce': False, 'overlap_param_gather': False, 'disable_last_saving': True, 'disable_wandb': True, 'disable_tensorboard': True, 'profile': False, 'use_pytorch_profiler': False, 'profile_step_end': 7, 'profile_step_start': 6} [Primus:Env] HF_HOME already set: /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/data/huggingface diff --git a/deepseek-v4/develop/progress/p22/smoke_eager_pp2_ep4.log b/deepseek-v4/develop/progress/p22/smoke_eager_pp2_ep4.log index 34b48b5e2..ebd895714 100644 --- a/deepseek-v4/develop/progress/p22/smoke_eager_pp2_ep4.log +++ b/deepseek-v4/develop/progress/p22/smoke_eager_pp2_ep4.log @@ -126,7 +126,7 @@ uv.lock' ']' [2026-05-08 02:05:56] [NODE-?(mi355-gpu-12)] [INFO] ----------------------------------------------- [2026-05-08 02:05:56] [NODE-?(mi355-gpu-12)] [INFO] primus-cli-direct.sh [2026-05-08 02:05:56] [NODE-?(mi355-gpu-12)] [INFO] ----------------------------------------------- -[2026-05-08 02:05:56] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRE_PARSE_ARGS (runner options): +[2026-05-08 02:05:56] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRE_PARSE_ARGS (runner options): [2026-05-08 02:05:56] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRIMUS_ARGS (python module args): train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 2 --expert_model_parallel_size 4 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --enable_primus_turbo False --use_turbo_attention False --use_turbo_deepep False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 [2026-05-08 02:05:56] [NODE-?(mi355-gpu-12)] [INFO] [direct] Combined args for second pass: -- train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 2 --expert_model_parallel_size 4 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --enable_primus_turbo False --use_turbo_attention False --use_turbo_deepep False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 [2026-05-08 02:05:56] [NODE-?(mi355-gpu-12)] [INFO] [direct] Loaded config: /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/runner/.primus.yaml @@ -331,11 +331,11 @@ WARNING: AMD GPU device(s) is/are in a low-power state. Check power control/runt Executing: torchrun --nproc_per_node 8 --nnodes 1 --node_rank 0 --master_addr localhost --master_port 1234 --local-ranks-filter 0 primus/cli/main.py train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 2 --expert_model_parallel_size 4 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --enable_primus_turbo False --use_turbo_attention False --use_turbo_deepep False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM 2>&1 | tee logs/log_20260508_020556.txt ========================================== - + ========================================== -W0508 02:05:59.345000 87316 torch/distributed/run.py:852] +W0508 02:05:59.345000 87316 torch/distributed/run.py:852] W0508 02:05:59.345000 87316 torch/distributed/run.py:852] ***************************************** -W0508 02:05:59.345000 87316 torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. +W0508 02:05:59.345000 87316 torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. W0508 02:05:59.345000 87316 torch/distributed/run.py:852] ***************************************** [Primus:Runtime] Applying CLI overrides: {'num_layers': 8, 'train_iters': 10, 'lr_warmup_iters': 0, 'lr_decay_iters': 10, 'micro_batch_size': 1, 'global_batch_size': 16, 'seq_length': 128, 'max_position_embeddings': 128, 'rope_type': 'rope', 'tensor_model_parallel_size': 1, 'pipeline_model_parallel_size': 2, 'expert_model_parallel_size': 4, 'num_experts': 8, 'moe_router_topk': 2, 'moe_router_enable_expert_bias': False, 'moe_ffn_hidden_size': 512, 'index_topk': 8, 'v4_grouped_experts_support_clamped_swiglu': True, 'compress_ratios': '[0,0,4,4,4,4,4,0]', 'mtp_num_layers': 0, 'mock_data': True, 'enable_primus_turbo': False, 'use_turbo_attention': False, 'use_turbo_deepep': False, 'use_turbo_grouped_mlp': False, 'moe_use_legacy_grouped_gemm': False, 'fp8': 'null', 'fp8_recipe': 'null', 'recompute_num_layers': 0, 'recompute_granularity': 'full', 'recompute_method': 'block', 'overlap_grad_reduce': False, 'overlap_param_gather': False, 'disable_last_saving': True, 'disable_wandb': True, 'disable_tensorboard': True, 'profile': False, 'use_pytorch_profiler': False, 'profile_step_end': 7, 'profile_step_start': 6} [Primus:Env] HF_HOME already set: /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/data/huggingface diff --git a/deepseek-v4/develop/progress/p22/smoke_turbo_pp2_ep4.log b/deepseek-v4/develop/progress/p22/smoke_turbo_pp2_ep4.log index 888760020..fd5447260 100644 --- a/deepseek-v4/develop/progress/p22/smoke_turbo_pp2_ep4.log +++ b/deepseek-v4/develop/progress/p22/smoke_turbo_pp2_ep4.log @@ -126,7 +126,7 @@ uv.lock' ']' [2026-05-08 02:12:24] [NODE-?(mi355-gpu-12)] [INFO] ----------------------------------------------- [2026-05-08 02:12:24] [NODE-?(mi355-gpu-12)] [INFO] primus-cli-direct.sh [2026-05-08 02:12:24] [NODE-?(mi355-gpu-12)] [INFO] ----------------------------------------------- -[2026-05-08 02:12:24] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRE_PARSE_ARGS (runner options): +[2026-05-08 02:12:24] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRE_PARSE_ARGS (runner options): [2026-05-08 02:12:24] [NODE-?(mi355-gpu-12)] [INFO] [direct] PRIMUS_ARGS (python module args): train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 2 --expert_model_parallel_size 4 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --enable_primus_turbo True --use_turbo_attention True --use_turbo_deepep False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 [2026-05-08 02:12:24] [NODE-?(mi355-gpu-12)] [INFO] [direct] Combined args for second pass: -- train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 2 --expert_model_parallel_size 4 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --enable_primus_turbo True --use_turbo_attention True --use_turbo_deepep False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 [2026-05-08 02:12:24] [NODE-?(mi355-gpu-12)] [INFO] [direct] Loaded config: /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/runner/.primus.yaml @@ -331,11 +331,11 @@ WARNING: AMD GPU device(s) is/are in a low-power state. Check power control/runt Executing: torchrun --nproc_per_node 8 --nnodes 1 --node_rank 0 --master_addr localhost --master_port 1234 --local-ranks-filter 0 primus/cli/main.py train pretrain --config examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM --num_layers 8 --train_iters 10 --lr_warmup_iters 0 --lr_decay_iters 10 --micro_batch_size 1 --global_batch_size 16 --seq_length 128 --max_position_embeddings 128 --rope_type rope --tensor_model_parallel_size 1 --pipeline_model_parallel_size 2 --expert_model_parallel_size 4 --num_experts 8 --moe_router_topk 2 --moe_router_enable_expert_bias False --moe_ffn_hidden_size 512 --index_topk 8 --v4_grouped_experts_support_clamped_swiglu True --compress_ratios [0,0,4,4,4,4,4,0] --mtp_num_layers 0 --mock_data True --enable_primus_turbo True --use_turbo_attention True --use_turbo_deepep False --use_turbo_grouped_mlp False --moe_use_legacy_grouped_gemm False --fp8 null --fp8_recipe null --recompute_num_layers 0 --recompute_granularity full --recompute_method block --overlap_grad_reduce False --overlap_param_gather False --disable_last_saving True --disable_wandb True --disable_tensorboard True --profile False --use_pytorch_profiler False --profile_step_end 7 --profile_step_start 6 --backend_path /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/third_party/Megatron-LM 2>&1 | tee logs/log_20260508_021224.txt ========================================== - + ========================================== -W0508 02:12:27.225000 92046 torch/distributed/run.py:852] +W0508 02:12:27.225000 92046 torch/distributed/run.py:852] W0508 02:12:27.225000 92046 torch/distributed/run.py:852] ***************************************** -W0508 02:12:27.225000 92046 torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. +W0508 02:12:27.225000 92046 torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. W0508 02:12:27.225000 92046 torch/distributed/run.py:852] ***************************************** [Primus:Runtime] Applying CLI overrides: {'num_layers': 8, 'train_iters': 10, 'lr_warmup_iters': 0, 'lr_decay_iters': 10, 'micro_batch_size': 1, 'global_batch_size': 16, 'seq_length': 128, 'max_position_embeddings': 128, 'rope_type': 'rope', 'tensor_model_parallel_size': 1, 'pipeline_model_parallel_size': 2, 'expert_model_parallel_size': 4, 'num_experts': 8, 'moe_router_topk': 2, 'moe_router_enable_expert_bias': False, 'moe_ffn_hidden_size': 512, 'index_topk': 8, 'v4_grouped_experts_support_clamped_swiglu': True, 'compress_ratios': '[0,0,4,4,4,4,4,0]', 'mtp_num_layers': 0, 'mock_data': True, 'enable_primus_turbo': True, 'use_turbo_attention': True, 'use_turbo_deepep': False, 'use_turbo_grouped_mlp': False, 'moe_use_legacy_grouped_gemm': False, 'fp8': 'null', 'fp8_recipe': 'null', 'recompute_num_layers': 0, 'recompute_granularity': 'full', 'recompute_method': 'block', 'overlap_grad_reduce': False, 'overlap_param_gather': False, 'disable_last_saving': True, 'disable_wandb': True, 'disable_tensorboard': True, 'profile': False, 'use_pytorch_profiler': False, 'profile_step_end': 7, 'profile_step_start': 6} [Primus:Env] HF_HOME already set: /shared/amdgpu/home/wen_xie_qle/workspace/Primus-deepseek-v4/data/huggingface @@ -2485,7 +2485,7 @@ Traceback (most recent call last): ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 317, in launch_agent raise ChildFailedError( -torch.distributed.elastic.multiprocessing.errors.ChildFailedError: +torch.distributed.elastic.multiprocessing.errors.ChildFailedError: ============================================================ primus/cli/main.py FAILED ------------------------------------------------------------ diff --git a/deepseek-v4/develop/progress/p35/bench_rope_triton.py b/deepseek-v4/develop/progress/p35/bench_rope_triton.py index a933e6a2d..75ccc8d57 100644 --- a/deepseek-v4/develop/progress/p35/bench_rope_triton.py +++ b/deepseek-v4/develop/progress/p35/bench_rope_triton.py @@ -45,7 +45,7 @@ if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) -from primus.backends.megatron.core.transformer.v4_attention_kernels._triton.rope_interleaved_partial import ( # noqa: E402 +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rope_interleaved_partial import ( # noqa: E402 RoPEInterleavedPartialFn, eager_apply_interleaved_partial_rope, ) diff --git a/deepseek-v4/develop/progress/p36/bench_sinkhorn.py b/deepseek-v4/develop/progress/p36/bench_sinkhorn.py index 8a880fde9..8f9e64b60 100644 --- a/deepseek-v4/develop/progress/p36/bench_sinkhorn.py +++ b/deepseek-v4/develop/progress/p36/bench_sinkhorn.py @@ -51,7 +51,7 @@ _get_compiled_sinkhorn, sinkhorn_normalize, ) -from primus.backends.megatron.core.transformer.v4_attention_kernels._triton.sinkhorn import ( # noqa: E402 +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.sinkhorn import ( # noqa: E402 SinkhornNormalizeFn, eager_sinkhorn_normalize, ) diff --git a/deepseek-v4/develop/progress/p37/bench_hc_glue.py b/deepseek-v4/develop/progress/p37/bench_hc_glue.py index aa64c1b66..da2682f37 100644 --- a/deepseek-v4/develop/progress/p37/bench_hc_glue.py +++ b/deepseek-v4/develop/progress/p37/bench_hc_glue.py @@ -38,7 +38,7 @@ if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) -from primus.backends.megatron.core.transformer.v4_attention_kernels._triton.hc_glue import ( # noqa: E402 +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_glue import ( # noqa: E402 HCComputeTailFn, ) diff --git a/deepseek-v4/develop/progress/p38/bench/small.json b/deepseek-v4/develop/progress/p38/bench/small.json index 57eee4131..93388d608 100644 --- a/deepseek-v4/develop/progress/p38/bench/small.json +++ b/deepseek-v4/develop/progress/p38/bench/small.json @@ -47,4 +47,4 @@ } } } -} \ No newline at end of file +} diff --git a/deepseek-v4/develop/progress/p38/bench/v4.json b/deepseek-v4/develop/progress/p38/bench/v4.json index ac5d58a3a..52f15c0da 100644 --- a/deepseek-v4/develop/progress/p38/bench/v4.json +++ b/deepseek-v4/develop/progress/p38/bench/v4.json @@ -47,4 +47,4 @@ } } } -} \ No newline at end of file +} diff --git a/deepseek-v4/develop/progress/p38/bench_indexer.py b/deepseek-v4/develop/progress/p38/bench_indexer.py index e5813c778..69dda60af 100644 --- a/deepseek-v4/develop/progress/p38/bench_indexer.py +++ b/deepseek-v4/develop/progress/p38/bench_indexer.py @@ -30,11 +30,10 @@ if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) -from primus.backends.megatron.core.transformer.v4_attention_kernels._triton.indexer_score import ( # noqa: E402 +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score import ( # noqa: E402 IndexerScoreFn, ) - _MODES = { "v4": dict(B=1, S=4096, P=1024, H=8, HD=128), "small": dict(B=2, S=128, P=32, H=8, HD=128), @@ -85,9 +84,17 @@ def _build_pool(args): pool = [] for c in range(max(1, args.n_input_copies)): gen = torch.Generator(device=device).manual_seed(args.seed + c) - q = torch.randn((B, S, H, HD), dtype=_dtype(args.dtype), device=device, generator=gen).requires_grad_(True) - k = torch.randn((B, P, HD), dtype=_dtype(args.dtype), device=device, generator=gen).requires_grad_(True) - w = torch.randn((B, S, H), dtype=_dtype(args.dtype), device=device, generator=gen).abs().requires_grad_(True) + q = torch.randn((B, S, H, HD), dtype=_dtype(args.dtype), device=device, generator=gen).requires_grad_( + True + ) + k = torch.randn((B, P, HD), dtype=_dtype(args.dtype), device=device, generator=gen).requires_grad_( + True + ) + w = ( + torch.randn((B, S, H), dtype=_dtype(args.dtype), device=device, generator=gen) + .abs() + .requires_grad_(True) + ) pool.append((q, k, w)) return pool, cfg @@ -179,8 +186,12 @@ def main() -> None: results = {} for name, triton_path in [("triton", True), ("eager", False)]: _time_fwd(pool, cfg, _dtype(args.dtype), triton_path=triton_path, iters=args.warmup, flusher=flusher) - fwd = _time_fwd(pool, cfg, _dtype(args.dtype), triton_path=triton_path, iters=args.iters, flusher=flusher) - bwd = _time_bwd(pool, cfg, _dtype(args.dtype), triton_path=triton_path, iters=args.iters, flusher=flusher) + fwd = _time_fwd( + pool, cfg, _dtype(args.dtype), triton_path=triton_path, iters=args.iters, flusher=flusher + ) + bwd = _time_bwd( + pool, cfg, _dtype(args.dtype), triton_path=triton_path, iters=args.iters, flusher=flusher + ) fwd_med = median(fwd) bwd_med = median(bwd) results[name] = { diff --git a/deepseek-v4/develop/progress/p41/bench_indexer_tail.py b/deepseek-v4/develop/progress/p41/bench_indexer_tail.py index d02240610..50d06de21 100644 --- a/deepseek-v4/develop/progress/p41/bench_indexer_tail.py +++ b/deepseek-v4/develop/progress/p41/bench_indexer_tail.py @@ -43,10 +43,10 @@ if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) -from primus.backends.megatron.core.transformer.v4_attention_kernels._triton.indexer_score import ( # noqa: E402 +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score import ( # noqa: E402 IndexerScoreFn, ) -from primus.backends.megatron.core.transformer.v4_attention_kernels._triton.indexer_score_post import ( # noqa: E402 +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score_post import ( # noqa: E402 IndexerScorePostFn, ) diff --git a/deepseek-v4/develop/progress/plan-2-summary.md b/deepseek-v4/develop/progress/plan-2-summary.md index eb5b2edf3..347317fa5 100644 --- a/deepseek-v4/develop/progress/plan-2-summary.md +++ b/deepseek-v4/develop/progress/plan-2-summary.md @@ -174,4 +174,3 @@ profile launchers + matching `.log` outputs). - Profile traces: see §3.2 above (rank-0 Chrome-trace JSON in the `output/amd/tas-mi355x-20260507/p19_profile_`* `tensorboard/` folders). - Day-1 baseline (plan-0 / plan-1): `[progress/2026-04-28-day-1-summary.md](./2026-04-28-day-1-summary.md)`. - diff --git a/deepseek-v4/projection/README.md b/deepseek-v4/projection/README.md index 700fd7e69..a6ae0c25d 100644 --- a/deepseek-v4/projection/README.md +++ b/deepseek-v4/projection/README.md @@ -28,11 +28,12 @@ run profiling script (per cr) -> chrome trace JSON (rank 0) breakdown JSON (site/data/.json) | v - static site (site/index.html) + 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 ``` @@ -46,6 +47,7 @@ deepseek-v4/projection/ 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 diff --git a/deepseek-v4/projection/design/01-overview.md b/deepseek-v4/projection/design/01-overview.md index 0b3de2487..eed63aac8 100644 --- a/deepseek-v4/projection/design/01-overview.md +++ b/deepseek-v4/projection/design/01-overview.md @@ -73,10 +73,11 @@ We capture with the distributed optimizer's comm overlap **off**. Two reasons: 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 **grouping kernels by `(module, phase, shape)` and -taking the per-group `min`** then removes residual run-to-run jitter and one-off -warm-up effects, yielding a stable clean time per kernel. Calls-per-microbatch = -launches / GA. +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 diff --git a/deepseek-v4/projection/design/02-assumptions.md b/deepseek-v4/projection/design/02-assumptions.md index 49b17896f..c86e751d8 100644 --- a/deepseek-v4/projection/design/02-assumptions.md +++ b/deepseek-v4/projection/design/02-assumptions.md @@ -17,8 +17,10 @@ start here. 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** (total params / DP under zero1 sharding) and is treated as - memory-bound (read/write params + states). + 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 @@ -52,8 +54,16 @@ start here. 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 is determined by the stack - (backward frames) and/or `_fwd_`/`_bwd_` in kernel names. + 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 @@ -65,10 +75,17 @@ start here. 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) -- **A17.** MTP is ignored (Pro paper has MTP depth 1; Primus V4 lacks it). It is a - known missing term in the full-model projection. - **A18.** No cross-node EP; no TP; no CP comm modeling. - **A19.** FP8 / MXFP8 not modeled; BF16 only. diff --git a/deepseek-v4/projection/design/04-projection-math.md b/deepseek-v4/projection/design/04-projection-math.md index ac272e7d2..fbe192bcf 100644 --- a/deepseek-v4/projection/design/04-projection-math.md +++ b/deepseek-v4/projection/design/04-projection-math.md @@ -27,22 +27,72 @@ 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 cr (site control: none / full / selective). v1 -exposes none | full; full sets it to 1 for all 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`. Total model chunks `C = PP * VPP`, layers per chunk -`Lc = num_layers / C` (assume divisible; otherwise distribute remainder to the -first chunks). Build the ordered layer list from `compress_ratios`, slice it into -`C` contiguous chunks (Megatron interleaved assigns chunk `k` to device -`k mod PP`). For device `d`: +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)] @@ -51,6 +101,7 @@ 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: ``` @@ -77,15 +128,19 @@ pipe_compute = (GA + (PP - 1) / VPP) * (Df_crit + Db_crit) Per-iteration, once, zero1-sharded, memory-bound: ``` -per_rank_opt_params = total_params / DP -opt_bytes = per_rank_opt_params * bytes_per_param * rw_factor +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` (dense + expert params; experts -are EP-sharded so per-rank expert params = experts/EP). `bytes_per_param`, -`rw_factor` (read+write ≈ 2), `opt_efficiency` are tunable; the measured -`optimizer.time_us` is used to calibrate `opt_time_per_param` for the MI355X -page and as a sanity check. +`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) @@ -102,6 +157,7 @@ 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 ``` diff --git a/deepseek-v4/projection/design/07-iteration-timeline.md b/deepseek-v4/projection/design/07-iteration-timeline.md new file mode 100644 index 000000000..1dcc38386 --- /dev/null +++ b/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 µs` (the op's + duration) and `starts @ 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/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh b/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh index 0ebb6e1e4..8b18eb930 100755 --- a/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh +++ b/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh @@ -20,8 +20,9 @@ # with overlap off every # microbatch's compute is # already clean) -# * GA = 2 => GBS = 2 * DP * MBS (min-grouping removes -# residual jitter) +# * 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) # @@ -89,8 +90,9 @@ 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 takes the per-kernel min over the two -# launches, recovering the overlap-free time (design/01-overview.md). +# 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) ---------------- @@ -109,8 +111,8 @@ 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; GA=2 + min-grouping then only removes residual -# jitter. (Production DP comm is assumed hidden in the projection anyway; A2.) +# 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 @@ -122,17 +124,25 @@ 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 Triton attn + Turbo MoE) -------------- +# ---------- 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_TRITON_ATTENTION=${USE_V4_TRITON_ATTENTION:-True} -export USE_V4_TRITON_CSA_ATTENTION=${USE_V4_TRITON_CSA_ATTENTION:-True} -export USE_V4_TILELANG_ATTENTION=${USE_V4_TILELANG_ATTENTION:-False} -export USE_V4_TILELANG_CSA_ATTENTION=${USE_V4_TILELANG_CSA_ATTENTION:-False} -export USE_V4_COMPILED_SINKHORN=${USE_V4_COMPILED_SINKHORN:-False} +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 @@ -149,11 +159,22 @@ if [ "$USE_TURBO_DEEPEP" = "True" ]; then fi export HSA_NO_SCRATCH_RECLAIM=${HSA_NO_SCRATCH_RECLAIM:-1} -# Wider profiler window (active 3 steps) so at least one full training iteration -# (fwd+bwd+opt) is captured cleanly; min-grouping in the parser then picks the -# clean step. A narrow 1-step window occasionally aligned to the optimizer tail -# only (cr0/cr128), missing the compute kernels. -export TRAIN_ITERS=${TRAIN_ITERS:-12} +# 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 @@ -206,10 +227,9 @@ mkdir -p "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" "${DISTOPT_ARGS[@]}" \ --enable_primus_turbo "$ENABLE_PRIMUS_TURBO" \ --use_turbo_attention "$USE_TURBO_ATTENTION" \ - --use_v4_triton_attention "$USE_V4_TRITON_ATTENTION" \ - --use_v4_triton_csa_attention "$USE_V4_TRITON_CSA_ATTENTION" \ - --use_v4_tilelang_attention "$USE_V4_TILELANG_ATTENTION" \ - --use_v4_tilelang_csa_attention "$USE_V4_TILELANG_CSA_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[@]}" \ @@ -218,6 +238,7 @@ mkdir -p "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" --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 \ @@ -226,6 +247,6 @@ mkdir -p "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" --profile True \ --use_pytorch_profiler True \ "${PROFILER_ACTIVITY_ARGS[@]}" \ - --profile_step_start 6 \ - --profile_step_end 9 \ + --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/deepseek-v4/projection/site/assets/app.js b/deepseek-v4/projection/site/assets/app.js index 30db8dafe..b1b0b32bd 100644 --- a/deepseek-v4/projection/site/assets/app.js +++ b/deepseek-v4/projection/site/assets/app.js @@ -8,6 +8,17 @@ 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); @@ -47,66 +58,123 @@ function defaultControls(data) { 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 : 512, - recompute: "full", - optEff: 0.7, computeEff: 1.0, calibFactor: 0.93, bytesPerParam: data.optimizer?.bytes_per_param || 18, + 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 = [ - ["world", "World size (GPUs)", "int"], - ["pp", "PP (pipeline)", "int"], - ["vpp", "VPP (interleave)", "int"], - ["ep", "EP (expert)", "int"], - ["dp", "DP (derived)", "ro"], - ["tp", "TP (tensor)", "int"], - ["cp", "CP (context)", "int"], - ["mbs", "Micro batch size", "int"], - ["gbs", "Global batch size", "int"], - ["recompute", "Recompute", "sel"], - ["bytesPerParam", "Optim bytes/param", "int"], - ["calibFactor", "Calibration factor", "f"], - ["optEff", "Optim efficiency", "f"], - ["computeEff", "MI455 compute eff", "f"], - ["peak355", "MI355 peak TFLOPs", "int"], - ["bw355", "MI355 HBM GB/s", "int"], - ["peak455", "MI455 peak TFLOPs", "int"], - ["bw455", "MI455 HBM GB/s", "int"], + { 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) => Math.max(1, Math.floor(c.world / (c.pp * c.tp * c.cp))); +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 [key, label, kind] of CONTROL_DEFS) { + 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"]) { + 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") { - input = el("input", { id: `ctl-${key}`, value: derivedDP(c), disabled: "true" }); + 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", () => { - const v = kind === "sel" ? input.value : Number(input.value); - c[key] = v; + c[key] = parseControlValue(input, kind); renderAll(); }); } @@ -115,6 +183,118 @@ function renderControls() { } } +// 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) // --------------------------------------------------------------------------- @@ -144,22 +324,206 @@ function layerTimes(data, cr, gpu, c) { let bwd = sumTime(aB, gpu, c) + sumTime(mB, gpu, c); let fFlops = sumFlops(aF) + sumFlops(mF); let bFlops = sumFlops(aB) + sumFlops(mB); - if (c.recompute === "full") { - bwd += fwd; // replay one forward (Step 1) - bFlops += fFlops; - } 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) // --------------------------------------------------------------------------- @@ -173,38 +537,103 @@ function estimateParams(cfg) { 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) { +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 = derivedDP(c); + const dp = validation.dp; const world = c.world; const lt = {}; - for (const cr of ["0", "4", "128"]) lt[cr] = layerTimes(data, cr, gpu, c); + 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); - for (let i = 0; i < L; i++) { - const chunk = Math.floor(i / perChunk); - const dev = chunk % c.pp; + 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] += t.bwd; + 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 - Df[0] += nonLayer(data, "embedding", "forward", gpu, c); - Db[0] += nonLayer(data, "embedding", "backward", gpu, c); + // 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; - Df[last] += nonLayer(data, "output", "forward", gpu, c) + nonLayer(data, "loss", "forward", gpu, c); - Db[last] += nonLayer(data, "output", "backward", gpu, c) + nonLayer(data, "loss", "backward", gpu, c); + 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). @@ -212,11 +641,12 @@ function project(data, gpu, c) { 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 world) - const totalParams = estimateParams(cfg); - const perRankParams = totalParams / world; + // 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 * 2) / bw / c.optEff; + const optTimeS = (perRankParams * c.bytesPerParam) / bw / c.optEff; const optUs = optTimeS * 1e6; // Step 5: totals @@ -236,6 +666,12 @@ function project(data, gpu, c) { 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"); @@ -245,7 +681,12 @@ function project(data, gpu, c) { return { lt, Df, Db, critF, critB, ga, pipeUs, bubbleFrac, world, dp, totalParams, - perRankParams, optUs, iterUs, tokIter, tokS, tokSgpu, flopsIter, tflopsGpu, seq, + 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, }; } @@ -266,6 +707,7 @@ function renderConfig() { ["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"]], @@ -329,8 +771,39 @@ function breakdownBlock(title, fwdList, bwdList) { 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; @@ -354,13 +827,32 @@ function step(label, value, sub) { return s; } -function renderResults() { +function renderResults(validation) { const c = STATE.controls, gpu = STATE.gpu, d = STATE.data; $("#results-gpu").textContent = `· ${gpu}`; - const p = project(d, gpu, c); + 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)); @@ -371,18 +863,25 @@ function renderResults() { head.append(mk("Iteration time", `${fmt(p.iterUs / 1000, 1)} ms`)); head.append(mk("Pipeline bubble", `${fmt(p.bubbleFrac * 100, 1)} %`)); - const steps = $("#results-steps"); - steps.innerHTML = ""; 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(" · "); - steps.append(step("Per-layer fwd/bwd (µs, " + (c.recompute === "full" ? "recompute on" : "no recompute") + ")", "", ltDesc)); + 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; per-device fwd=[${p.Df.map((x) => fmt(x / 1000, 1)).join(", ")}] ms`)); - steps.append(step("GA (microbatches)", fmtInt(p.ga), `GA = GBS ${c.gbs} / (DP ${p.dp} × MBS ${c.mbs})`)); + `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 params/rank × ${c.bytesPerParam}B ×2 / HBM-BW / eff ${c.optEff}`)); + `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}`)); @@ -397,14 +896,589 @@ function renderResults() { } } +// --------------------------------------------------------------------------- +// 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 + // (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("]*>)/, `$1`); + 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: + 'attn' + + 'mlp (experts)' + + 'a2a (dispatch+combine)' + + 'misc/unattributed' }); +} + +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: + 'manual per-layer total (no module split)' })); + 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: + 'cr=0cr=4cr=128' + + 'embeddingoutput/loss/MTP' })); +} + +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: + 'forwardbackward' + + (sim.VPP > 1 ? 'lighter→darker = VPP chunk 0→' + (sim.VPP - 1) + '' : '') + + 'gaps = pipeline bubble (idle)' })); + 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 = derivedDP(w); + if (roDp) roDp.value = Number.isFinite(validation.dp) ? validation.dp : "—"; + renderValidation(validation); + renderModeSwitch(); + renderManualGrid(); renderConfig(); renderBreakdown(); - renderResults(); + renderResults(validation); + renderTimeline(); } // --------------------------------------------------------------------------- @@ -425,22 +1499,69 @@ async function init(model) { } } -$("#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); -}); +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(".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; - renderBreakdown(); - renderResults(); + document.querySelectorAll(".tl-view").forEach((tab) => { + tab.addEventListener("click", () => { + STATE.tlView = tab.dataset.view; + renderTimeline(); + }); }); -}); -init(modelFromQuery()); + // 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/deepseek-v4/projection/site/assets/style.css b/deepseek-v4/projection/site/assets/style.css index 2f2c9aa5f..8d9d9d805 100644 --- a/deepseek-v4/projection/site/assets/style.css +++ b/deepseek-v4/projection/site/assets/style.css @@ -27,7 +27,7 @@ body { border-bottom: 1px solid var(--border); padding: 28px 0; } -.hero__inner, .layout { max-width: 1200px; margin: 0 auto; padding: 0 24px; } +.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; } @@ -40,7 +40,44 @@ body { } .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; } @@ -55,6 +92,15 @@ 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); } @@ -82,6 +128,45 @@ table.bd td.rowlab, table.bd th.rowlab { text-align: left; color: var(--muted); .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 @@ -89,21 +174,160 @@ table.bd td.rowlab, table.bd th.rowlab { text-align: left; color: var(--muted); .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: 1200px; margin: 0 auto; } +.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; } } +@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/deepseek-v4/projection/site/data/flash.json b/deepseek-v4/projection/site/data/flash.json index 43f5d611f..31a744e44 100644 --- a/deepseek-v4/projection/site/data/flash.json +++ b/deepseek-v4/projection/site/data/flash.json @@ -8,18 +8,31 @@ "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 (B=1) Megatron-convention FLOPs at capture seq; site multiplies by cr_layer_counts x GA x DP" + "note": "per-layer and MTP (B=1) Megatron-convention FLOPs at capture seq; site multiplies by GA x DP" }, - "generated_at": "2026-06-22T17:52:01Z", + "generated_at": "2026-06-30T13:59:46Z", "provenance": { "traces": { - "0": "/apps/tas/wenx/workspace/Primus-deepseek-v4/output/amd/tas-mi355x-20260618/projection_flash_cr0_seq4096_ep8/tensorboard/primus-megatron-exp[projection_flash_cr0_seq4096_ep8]-rank[0].1781793844365290266.pt.trace.json", - "4": "/apps/tas/wenx/workspace/Primus-deepseek-v4/output/amd/tas-mi355x-20260618/projection_flash_cr4_seq4096_ep8/tensorboard/primus-megatron-exp[projection_flash_cr4_seq4096_ep8]-rank[0].1781793915830474166.pt.trace.json", - "128": "/apps/tas/wenx/workspace/Primus-deepseek-v4/output/amd/tas-mi355x-20260618/projection_flash_cr128_seq4096_ep8/tensorboard/primus-megatron-exp[projection_flash_cr128_seq4096_ep8]-rank[0].1781793986369245582.pt.trace.json" + "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": [], - "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": { + "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", @@ -31,7 +44,15 @@ "optimizer": "adam", "distributed_optimizer": false, "recompute": "off", - "measured_iter_time_ms": null + "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, @@ -44,6 +65,11 @@ "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, @@ -94,7 +120,7 @@ "4": 20, "128": 20 }, - "total_params": 283810725888 + "total_params": 290419900416 }, "hardware": { "MI355X": { @@ -111,8 +137,8 @@ "attention": { "forward": [ { - "module": "other", - "time_us": 12248.881, + "module": "attn.misc", + "time_us": 2768.835, "class": "memory_bound", "flop_class": null, "flops": null, @@ -120,101 +146,101 @@ "kernels": [ { "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern", - "time_us": 6205.904 + "time_us": 600.524 }, { - "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel", - "time_us": 1238.199 + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::AUnaryFunctor", + "time_us": 252.844 }, { - "name": "void at::native::vectorized_templated_elementwise_kernel<4, at::native::CUDAFunc", - "time_us": 881.949 + "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::BinaryFunctor(std::bfl", - "time_us": 115.8 + "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", + "time_us": 88.64 }, { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern", - "time_us": 24.463 + "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::(anonymous", + "time_us": 29.89 }, { - "name": "void primus_turbo::compute_grouped_gemm_args", - "time_us": 14.43 - }, - { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_", - "time_us": 10.909 + "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024>(HIP_vector_type(int const*, int*, int ", + "time_us": 111.084 }, { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::FillFunctor", - "time_us": 1.659 + "name": "void primus_turbo::deep_ep::layout::get_dispatch_layout<256, 4, 8>(long const*, ", + "time_us": 19.553 } ] - } - ], - "backward": [ + }, { "module": "moe.combine", - "time_us": 1399.813, + "time_us": 321.999, "class": "memory_bound", "flop_class": null, "flops": null, "tflops": null, "kernels": [ { - "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", - "time_us": 937.017 + "name": "void primus_turbo::deep_ep::intranode::combine(hip_bfloat", + "time_us": 231.632 }, { - "name": "void primus_turbo::deep_ep::intranode::combine(hip_", - "time_us": 462.796 + "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", + "time_us": 90.367 } ] }, { - "module": "moe.dispatch", - "time_us": 618.803, - "class": "memory_bound", - "flop_class": null, - "flops": null, - "tflops": null, + "module": "moe.shared_expert", + "time_us": 112.521, + "class": "compute_bound", + "flop_class": "gemm", + "flops": 68719476736.0, + "tflops": 610.7, "kernels": [ { - "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024, true>(HIP_vector_type(int const*, int", - "time_us": 72.124 + "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", + "time_us": 21.219 }, { - "name": "void primus_turbo::deep_ep::intranode::notify_dispatch<8>(int const*, int*, int ", - "time_us": 72.07 + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa", + "time_us": 5.396 }, { - "name": "void primus_turbo::deep_ep::layout::get_dispatch_layout<256, 4, 8>(long const*, ", - "time_us": 19.856 + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor", + "time_us": 8.02 }, { - "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa", - "time_us": 21.546 + "name": "void at::native::_scatter_gather_elementwise_kernel<256, 4, at::native::_cuda_sc", + "time_us": 6.869 }, { - "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor(std::bfl", + "time_us": 115.553 }, { - "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor(hip_bfloat", + "time_us": 238.449 }, { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa", - "time_us": 5.506 - }, + "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 at::native::reduce_kernel<512, 1, at::native::ReduceOp(HIP_vector_type(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": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp", + "time_us": 247.801 }, { "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", - "time_us": 911.084 + "time_us": 192.373 }, { - "name": "void at::native::vectorized_templated_elementwise_kernel<4, at::native::CUDAFunc", - "time_us": 879.542 + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa", + "time_us": 150.529 }, { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::CUDAFunctor_add(hip_bfloat", + "time_us": 230.235 }, { - "name": "void at::native::(anonymous namespace)::CatArrayBatchedCopy(void**, int*, in", + "time_us": 100.447 } ] }, { - "module": "attn.norm", - "time_us": 310.559, + "module": "moe.dispatch", + "time_us": 279.761, "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": 107.056 - }, - { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa", - "time_us": 89.503 - }, - { - "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp(HIP_vector_type(int const*, int*, int ", + "time_us": 27.653 }, { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::cos_kernel_cuda(at", - "time_us": 6.152 + "name": "void primus_turbo::deep_ep::layout::get_dispatch_layout<256, 4, 8>(long const*, ", + "time_us": 19.776 } ] - } - ] - }, - "moe": { - "forward": [ + }, { - "module": "moe.grouped_gemm", - "time_us": 2887.082, + "module": "moe.shared_expert", + "time_us": 112.854, "class": "compute_bound", - "flop_class": "grouped_gemm", - "flops": null, - "tflops": null, + "flop_class": "gemm", + "flops": 68719476736.0, + "tflops": 608.9, "kernels": [ { - "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel(std::bfl", - "time_us": 117.143 + "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<4, at::native::bfloat16_copy_kern", - "time_us": 25.293 + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa", + "time_us": 5.393 }, { - "name": "void primus_turbo::compute_grouped_gemm_args", - "time_us": 14.103 + "time_us": 8.086 }, { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_", - "time_us": 10.496 + "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": 3.5 + "time_us": 5.669 }, { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::FillFunctor", - "time_us": 1.573 + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa", + "time_us": 4.696 } ] } ], "backward": [ { - "module": "moe.combine", - "time_us": 1375.809, - "class": "memory_bound", - "flop_class": null, + "module": "moe.grouped_gemm", + "time_us": 1896.364, + "class": "compute_bound", + "flop_class": "grouped_gemm", "flops": null, "tflops": null, "kernels": [ { - "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", - "time_us": 913.226 + "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel(hip_", - "time_us": 462.583 + "name": "void primus_turbo::grouped_gemm_variable_k_postprocess(std::bfl", + "time_us": 114.644 + }, + { + "name": "void primus_turbo::compute_grouped_gemm_variable_k_args(HIP_vector_type(HIP_vector_type(int const*, int", - "time_us": 80.144 - }, - { - "name": "void primus_turbo::deep_ep::layout::get_dispatch_layout<256, 4, 8>(long const*, ", - "time_us": 20.083 - }, - { - "name": "void primus_turbo::deep_ep::intranode::notify_dispatch<8>(int const*, int*, int ", - "time_us": 15.583 + "time_us": 178.088 } ] }, { - "module": "moe.grouped_gemm", - "time_us": 489.055, + "module": "moe.combine", + "time_us": 357.633, "class": "memory_bound", "flop_class": null, "flops": null, "tflops": null, "kernels": [ { - "name": "_stack_grouped_weight_bwd_kernel", - "time_us": 311.161 - }, - { - "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", - "time_us": 85.576 - }, - { - "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::(anonymous", - "time_us": 29.016 - }, - { - "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa", - "time_us": 21.916 - }, - { - "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor(hip_bfloat", + "time_us": 237.535 }, { - "name": "void at::native::vectorized_elementwise_kernel<8, at::native::CUDAFunctorOnSelf_", - "time_us": 15.55 + "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", + "time_us": 120.097 } ] }, { - "module": "moe.shared_expert", - "time_us": 113.441, - "class": "compute_bound", - "flop_class": "gemm", - "flops": 68719476736.0, - "tflops": 605.8, + "module": "moe.router", + "time_us": 44.726, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, "kernels": [ { - "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", - "time_us": 78.956 - }, - { - "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", - "time_us": 21.366 - }, - { - "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa", - "time_us": 5.413 - }, - { - "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor", + "time_us": 248.877 }, { - "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp(std::bfl", - "time_us": 117.196 + "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", + "time_us": 88.707 }, { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern", - "time_us": 23.84 + "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::(anonymous", + "time_us": 30.373 }, { - "name": "void primus_turbo::compute_grouped_gemm_args(HIP_vector_type", - "time_us": 14.28 - }, - { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_", - "time_us": 10.682 - }, - { - "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel", - "time_us": 3.5 + "name": "void primus_turbo::deep_ep::intranode::notify_dispatch<8>(int const*, int*, int ", + "time_us": 60.447 }, { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::FillFunctor", - "time_us": 1.579 + "name": "void primus_turbo::deep_ep::layout::get_dispatch_layout<256, 4, 8>(long const*, ", + "time_us": 20.186 } ] - } - ], - "backward": [ + }, { "module": "moe.combine", - "time_us": 1273.128, + "time_us": 309.929, "class": "memory_bound", "flop_class": null, "flops": null, "tflops": null, "kernels": [ { - "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", - "time_us": 810.992 + "name": "void primus_turbo::deep_ep::intranode::combine(hip_bfloat", + "time_us": 227.382 }, { - "name": "void primus_turbo::deep_ep::intranode::combine(hip_", - "time_us": 462.136 + "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", + "time_us": 82.547 } ] }, { - "module": "moe.dispatch", - "time_us": 845.998, - "class": "memory_bound", - "flop_class": null, - "flops": null, - "tflops": null, + "module": "moe.shared_expert", + "time_us": 112.514, + "class": "compute_bound", + "flop_class": "gemm", + "flops": 68719476736.0, + "tflops": 610.8, "kernels": [ { - "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024, true>(HIP_vector_type(int const*, int", - "time_us": 215.298 + "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", + "time_us": 21.266 }, { - "name": "void primus_turbo::deep_ep::intranode::notify_dispatch<8>(int const*, int*, int ", - "time_us": 154.151 + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa", + "time_us": 5.43 }, { - "name": "void primus_turbo::deep_ep::layout::get_dispatch_layout<256, 4, 8>(long const*, ", - "time_us": 20.68 + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor", + "time_us": 14.203 }, { - "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::(anonymous", - "time_us": 29.836 + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_", + "time_us": 10.482 }, { - "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa", - "time_us": 21.963 + "name": "void at::native::_scatter_gather_elementwise_kernel<256, 4, at::native::_cuda_sc", + "time_us": 7.482 }, { - "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor(std::bfl", + "time_us": 113.277 }, { - "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor(hip_bfloat", + "time_us": 233.195 }, { - "name": "void at::native::_scatter_gather_elementwise_kernel<256, 4, at::native::_cuda_sc", - "time_us": 7.512 - }, - { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa", - "time_us": 5.326 - }, + "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 at::native::reduce_kernel<512, 1, at::native::ReduceOp(HIP_vector_type(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": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp(char*, char*, long*, int, lo", - "time_us": 8.229 + "time_us": 7.84 } ] } @@ -1665,7 +1615,7 @@ "backward": [ { "module": "embedding", - "time_us": 88.911, + "time_us": 86.591, "class": "memory_bound", "flop_class": null, "flops": null, @@ -1673,27 +1623,27 @@ "kernels": [ { "name": "void at::native::(anonymous namespace)::compute_grad_weight", - "time_us": 34.727 + "time_us": 33.993 }, { - "name": "void at::native::(anonymous namespace)::sum_and_scatter(lon", - "time_us": 22.716 + "name": "void rocprim::ROCPRIM_400200_NS::detail::trampoline_kernel(lon", + "time_us": 21.21 }, { "name": "void rocprim::ROCPRIM_400200_NS::detail::init_lookback_scan_state_kernel(long*, ", - "time_us": 2.066 + "time_us": 2.08 }, { "name": "void at::native::(anonymous namespace)::krn_partials_per_segment(long*, lo", - "time_us": 2.029 + "time_us": 2.023 } ] } @@ -1703,58 +1653,58 @@ "forward": [ { "module": "output", - "time_us": 1603.856, + "time_us": 1553.935, "class": "compute_bound", "flop_class": "gemm", - "flops": 2171105968128.0, - "tflops": 1353.7, + "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": 1603.856 + "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": [ + "backward": [] + }, + "loss": { + "forward": [ { - "module": "output", - "time_us": 24.207, - "class": "compute_bound", - "flop_class": "gemm", - "flops": 268435456.0, - "tflops": 11.1, + "module": "loss", + "time_us": 368.303, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, "kernels": [ { - "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT16x16x512_MI16x16x1_SN_LDSB1_AFC1_AF", - "time_us": 24.207 + "name": "cross_entropy_kernel", + "time_us": 235.982 + }, + { + "name": "online_softmax_kernel", + "time_us": 132.321 } ] } - ] - }, - "loss": { - "forward": [], + ], "backward": [ { "module": "loss", - "time_us": 596.633, + "time_us": 197.392, "class": "memory_bound", "flop_class": null, "flops": null, "tflops": null, "kernels": [ - { - "name": "cross_entropy_kernel", - "time_us": 253.851 - }, { "name": "element_mul_kernel", - "time_us": 193.371 - }, - { - "name": "online_softmax_kernel", - "time_us": 149.411 + "time_us": 197.392 } ] } @@ -1764,10 +1714,10 @@ "optimizer": { "type": "adam", "measured_params": null, - "time_us": 14497.3, - "bytes_per_param": 18, + "time_us": 18459.0, + "bytes_per_param": 30, "class": "memory_bound", - "note": "measured one-layer optimizer-step kernel time on this rank; the site scales by per-rank params" + "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, diff --git a/deepseek-v4/projection/site/data/pro.json b/deepseek-v4/projection/site/data/pro.json index 1ab0c9d19..8cfd34cb4 100644 --- a/deepseek-v4/projection/site/data/pro.json +++ b/deepseek-v4/projection/site/data/pro.json @@ -8,10 +8,18 @@ "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 (B=1) Megatron-convention FLOPs at capture seq; site multiplies by cr_layer_counts x GA x DP" + "note": "per-layer and MTP (B=1) Megatron-convention FLOPs at capture seq; site multiplies by GA x DP" }, - "generated_at": "2026-06-22T17:51:57Z", + "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", @@ -19,7 +27,12 @@ "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": [], - "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": { + "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", @@ -31,7 +44,8 @@ "optimizer": "adam", "distributed_optimizer": false, "recompute": "off", - "measured_iter_time_ms": null + "measured_iter_time_ms": null, + "measured_anchor": null }, "model_config": { "num_layers": 61, @@ -44,6 +58,11 @@ "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, @@ -112,7 +131,7 @@ "4": 29, "128": 31 }, - "total_params": 1571740581888 + "total_params": 1597579198464 }, "hardware": { "MI355X": { @@ -129,99 +148,89 @@ "attention": { "forward": [ { - "module": "other", - "time_us": 50996.046, + "module": "attn.misc", + "time_us": 7402.198, "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": 38789.717 + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern", + "time_us": 1437.355 }, { - "name": "void at::native::vectorized_templated_elementwise_kernel<4, at::native::CUDAFunc", - "time_us": 3024.544 + "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp", + "time_us": 709.616 }, { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::BinaryFunctor", - "time_us": 726.343 + "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", + "time_us": 324.583 } ] }, { "module": "attn.proj", - "time_us": 5781.21, + "time_us": 1623.031, "class": "compute_bound", "flop_class": "gemm", - "flops": 6172539092992.0, - "tflops": 1067.7, + "flops": 1191719206912.0, + "tflops": 734.3, "kernels": [ - { - "name": "Cijk_Ailk_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", - "time_us": 3769.201 - }, { "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", - "time_us": 654.028 + "time_us": 905.959 }, { - "name": "Cijk_Ailk_Bjlk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x16x32_MI16x16x1_SN_LDSB1_AFC1", - "time_us": 489.093 + "name": "Cijk_Ailk_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x16x32_MI16x16x1_SN_LDSB1_AFC1", + "time_us": 476.296 }, { - "name": "Cijk_Ailk_Bjlk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", - "time_us": 201.328 + "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT32x64x128_MI16x16x1_SN_LDSB1_AFC1_AF", + "time_us": 117.74 }, { - "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_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_MT16x16x512_MI16x16x1_SN_LDSB0_AFC1", - "time_us": 161.16 - } - ] - }, - { - "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": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x128x64_MI16x16x1_SN_LDSB1_AFC", + "time_us": 21.693 }, { - "name": "_hc_compute_tail_fwd_kernel", - "time_us": 4.396 + "name": "Cijk_SS_BiasS_HAS_ScaleAlphaVec_PostGSU4_VW4", + "time_us": 4.239 } ] }, { "module": "attn.norm", - "time_us": 449.337, + "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 @@ -231,90 +240,100 @@ "time_us": 127.904 }, { - "name": "_apply_rope_fwd_kernel", - "time_us": 105.693 + "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp(std::bfl", - "time_us": 423.883 + "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 primus_turbo::compute_grouped_gemm_args", - "time_us": 14.06 - }, - { - "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel", - "time_us": 3.576 - }, - { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::FillFunctor", - "time_us": 2.073 - } - ] - } - ], - "backward": [ - { - "module": "moe.grouped_gemm", - "time_us": 1339.59, - "class": "memory_bound", - "flop_class": null, - "flops": null, - "tflops": null, - "kernels": [ - { - "name": "_stack_grouped_weight_bwd_kernel", - "time_us": 1055.078 - }, - { - "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", - "time_us": 135.577 - }, { "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::(anonymous", "time_us": 43.473 @@ -471,20 +412,12 @@ { "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor(hip_", - "time_us": 781.249 + "time_us": 386.653 }, { "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", - "time_us": 233.451 + "time_us": 41.18 } ] }, { "module": "moe.dispatch", - "time_us": 872.042, + "time_us": 422.925, "class": "memory_bound", "flop_class": null, "flops": null, @@ -510,11 +443,7 @@ "kernels": [ { "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024, true>(HIP_vector_type(int const*, int", - "time_us": 73.65 + "time_us": 382.809 }, { "name": "void primus_turbo::deep_ep::intranode::notify_dispatch<8>(int const*, int*, int ", @@ -562,109 +491,157 @@ }, { "module": "moe.router", - "time_us": 74.504, + "time_us": 91.185, "class": "memory_bound", "flop_class": null, "flops": null, "tflops": null, "kernels": [ { - "name": "_sinkhorn_bwd_kernel", - "time_us": 44.656 + "name": "_sinkhorn_fwd_kernel", + "time_us": 25.516 }, { - "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::bfloat16tofloat32_", + "time_us": 16.113 }, { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa", - "time_us": 6.666 + "name": "void at::native::vectorized_elementwise_kernel<16, at::native::FillFunctor", + "time_us": 14.06 }, { - "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp(std::bfl", + "time_us": 423.883 }, { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::BinaryFunctor", - "time_us": 993.69 + "name": "void primus_turbo::compute_grouped_gemm_args(hip_", + "time_us": 394.596 }, { - "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", - "time_us": 656.014 + "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(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": "Cijk_Ailk_Bjlk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", - "time_us": 200.238 + "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp", + "time_us": 709.102 + }, + { + "name": "void at::native::unrolled_elementwise_kernel(HIP_vector_type(int const*, int*, int ", + "time_us": 82.547 }, { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::compare_scalar_ker", - "time_us": 22.782 - }, + "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 at::native::elementwise_kernel_manual_unroll<128, 4, at::native::(anonymous", - "time_us": 6.495 + "name": "void primus_turbo::deep_ep::intranode::combine(hip_", + "time_us": 385.316 }, { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::cos_kernel_cuda(at", - "time_us": 6.475 + "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", + "time_us": 84.724 } ] - } - ] - }, - "moe": { - "forward": [ + }, { - "module": "moe.grouped_gemm", - "time_us": 8765.885, + "module": "moe.shared_expert", + "time_us": 254.469, "class": "compute_bound", - "flop_class": "grouped_gemm", + "flop_class": "gemm", "flops": null, "tflops": null, "kernels": [ { - "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel(std::bfl", - "time_us": 425.279 + "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<4, at::native::bfloat16_copy_kern", - "time_us": 43.96 + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa", + "time_us": 6.96 }, { - "name": "void primus_turbo::compute_grouped_gemm_args", - "time_us": 2.116 + "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 } ] } @@ -998,41 +1033,37 @@ "backward": [ { "module": "moe.grouped_gemm", - "time_us": 1356.649, - "class": "memory_bound", - "flop_class": null, + "time_us": 5908.193, + "class": "compute_bound", + "flop_class": "grouped_gemm", "flops": null, "tflops": null, "kernels": [ { - "name": "_stack_grouped_weight_bwd_kernel", - "time_us": 1074.128 - }, - { - "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", - "time_us": 135.8 + "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel(std::bfl", + "time_us": 425.279 }, { - "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa", - "time_us": 33.813 + "name": "void primus_turbo::compute_grouped_gemm_variable_k_args(HIP_vector_type(int const*, int", "time_us": 204.325 - }, - { - "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": 973.473, + "time_us": 503.433, "class": "memory_bound", "flop_class": null, "flops": null, @@ -1066,51 +1089,17 @@ "kernels": [ { "name": "void primus_turbo::deep_ep::intranode::combine(hip_", - "time_us": 779.425 + "time_us": 394.109 }, { "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", - "time_us": 194.047 - } - ] - }, - { - "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", + "time_us": 708.356 }, { - "name": "void at::native::vectorized_elementwise_kernel<4, at::native::BinaryFunctor", - "time_us": 730.665 + "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", + "time_us": 327.05 } ] }, { "module": "attn.proj", - "time_us": 5909.199, + "time_us": 1639.461, "class": "compute_bound", "flop_class": "gemm", - "flops": 6232668635136.0, - "tflops": 1054.7, + "flops": 1191719206912.0, + "tflops": 726.9, "kernels": [ { - "name": "Cijk_Ailk_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", - "time_us": 3794.003 + "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", + "time_us": 919.326 }, { - "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", - "time_us": 658.211 + "name": "Cijk_Ailk_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x16x32_MI16x16x1_SN_LDSB1_AFC1", + "time_us": 478.756 }, { - "name": "Cijk_Ailk_Bjlk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x16x32_MI16x16x1_SN_LDSB1_AFC1", - "time_us": 487.776 + "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT32x64x128_MI16x16x1_SN_LDSB1_AFC1_AF", + "time_us": 117.167 }, { - "name": "Cijk_Ailk_Bjlk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", - "time_us": 201.361 + "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT64x32x128_MI16x16x1_SN_LDSB1_AFC1_AF", + "time_us": 94.187 }, { - "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_MT256x128x64_MI16x16x1_SN_LDSB1_AFC", + "time_us": 21.563 }, { - "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT16x16x512_MI16x16x1_SN_LDSB0_AFC1", - "time_us": 160.704 + "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(hip_", + "time_us": 387.153 }, { - "name": "Cijk_SB_BiasS_HAS_ScaleAlphaVec_PostGSU4_VW4", - "time_us": 8.172 - }, + "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 at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", - "time_us": 5.979 + "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024, true>(HIP_vector_type(int const*, int*, int ", + "time_us": 45.283 }, { - "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp(long const*, ", + "time_us": 19.383 } ] - } - ] - }, - "moe": { - "forward": [ + }, { - "module": "moe.grouped_gemm", - "time_us": 8717.575, + "module": "moe.shared_expert", + "time_us": 254.795, "class": "compute_bound", - "flop_class": "grouped_gemm", + "flop_class": "gemm", "flops": null, "tflops": null, "kernels": [ { - "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel(std::bfl", - "time_us": 423.59 + "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<4, at::native::bfloat16_copy_kern", - "time_us": 44.276 + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa", + "time_us": 7.026 }, { - "name": "void primus_turbo::compute_grouped_gemm_args", - "time_us": 2.07 + "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 } ] } @@ -1526,41 +1535,37 @@ "backward": [ { "module": "moe.grouped_gemm", - "time_us": 1359.113, - "class": "memory_bound", - "flop_class": null, + "time_us": 5906.347, + "class": "compute_bound", + "flop_class": "grouped_gemm", "flops": null, "tflops": null, "kernels": [ { - "name": "_stack_grouped_weight_bwd_kernel", - "time_us": 1076.558 - }, - { - "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", - "time_us": 135.587 + "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel(std::bfl", + "time_us": 423.59 }, { - "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa", - "time_us": 33.963 + "name": "void primus_turbo::compute_grouped_gemm_variable_k_args(hip_", - "time_us": 781.462 + "time_us": 394.309 }, { "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", - "time_us": 212.091 + "time_us": 84.787 } ] }, { "module": "moe.dispatch", - "time_us": 857.885, + "time_us": 408.969, "class": "memory_bound", "flop_class": null, "flops": null, @@ -1586,59 +1591,17 @@ "kernels": [ { "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024, true>(HIP_vector_type(int const*, int*, int ", - "time_us": 45.283 + "time_us": 375.616 }, { "name": "void primus_turbo::deep_ep::intranode::cached_notify_dispatch<8>(int const*, int", "time_us": 33.353 - }, - { - "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::BinaryFunctorTraining Performance Projection multi-GPU projection. Page 1 is measured-MI355X; page 2 scales to MI455X by theoretical hardware ratios.

-
- - -
+ + + +
+
+ +
- +
+ +
+ -
+
@@ -51,27 +77,47 @@

Per-layer breakdown

(gemm / grouped_gemm / attn only); the headline TFLOP/s/GPU below uses the V4 analytic model FLOPs instead.

+
-
-
-

Projection controls

-
+
+

Projected throughput

+
+
+
+ +
+
+

Iteration timeline

- 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). + 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 design/07-iteration-timeline.md.

-
- -
-

Projected throughput

-
-
-
+
+
+ Layout +
+ + +
+
+
+ + + +
+
+
+
+
diff --git a/deepseek-v4/projection/tools/kernel_module_map.py b/deepseek-v4/projection/tools/kernel_module_map.py index fb442402d..fd6ddbcea 100644 --- a/deepseek-v4/projection/tools/kernel_module_map.py +++ b/deepseek-v4/projection/tools/kernel_module_map.py @@ -73,11 +73,16 @@ ("_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"), diff --git a/deepseek-v4/projection/tools/parse_trace.py b/deepseek-v4/projection/tools/parse_trace.py index 9d6e338f8..dad6291f7 100755 --- a/deepseek-v4/projection/tools/parse_trace.py +++ b/deepseek-v4/projection/tools/parse_trace.py @@ -24,14 +24,29 @@ 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 -from kernel_module_map import flop_class_from_kernel -from v4_flops import FB_FMA, layer_fmac, model_total_params, nonlayer_fmac +# 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] @@ -48,6 +63,9 @@ "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": { @@ -61,6 +79,9 @@ "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, }, } @@ -73,6 +94,26 @@ "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", @@ -107,6 +148,35 @@ def _dims_key(dims: Any) -> str: 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 @@ -139,6 +209,15 @@ def _resolve_module(kname: str, cpu_name: str, anc_classes: set[str], flop_class 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__" @@ -201,7 +280,8 @@ def has(*xs: str) -> bool: return "attn.proj" if flop_class == "gemm" else "attn.norm" if has("Embedding"): return "output" if flop_class == "gemm" else "embedding" - return "other" + fallback = coarse(module_from_kernel(kname)) + return fallback if fallback != "other" else "attn.misc" def parse_trace(path: Path, ga: int = 2): @@ -212,6 +292,9 @@ def parse_trace(path: Path, ga: int = 2): 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) @@ -224,6 +307,10 @@ def parse_trace(path: Path, ga: int = 2): 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: @@ -266,12 +353,58 @@ def enclosing(cpu: dict) -> set[str]: # 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 ev in kernels: + for i, ev in enumerate(kernels): name = ev.get("name", "") dur = float(ev["dur"]) ext = _arg(ev, "External id") @@ -286,16 +419,12 @@ def enclosing(cpu: dict) -> set[str]: if module == "__dpcomm__": dpcomm_us += dur continue - # Phase: a kernel-name _fwd_/_bwd_ tag is authoritative (V4 triton - # kernels encode it). Otherwise use the backward-only "Fwd thread id". - ln = name.lower() - if "_fwd" in ln and "_bwd" not in ln: - phase = "forward" - elif "_bwd" in ln: - phase = "backward" - else: - fwd_tid = _arg(cpu, "Fwd thread id") if cpu else None - phase = "backward" if (fwd_tid is not None or "ackward" in cpu_name) else "forward" + # 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( @@ -306,6 +435,7 @@ def enclosing(cpu: dict) -> set[str]: 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(): @@ -344,11 +474,12 @@ def enclosing(cpu: dict) -> set[str]: "tflops": round(tflops, 1) if tflops else None, "kernels": kern_list, } - # Logical bucket. attn.* and the unattributed `other` (dominated by the - # attention-side hyper-connection mixer / Indexer scalar ops) go to the - # attention bucket; this keeps the MoE bucket strictly the cr-independent - # router/dispatch/grouped_gemm/combine/shared_expert set (A10). - if module.startswith("attn.") or module == "other": + # 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" @@ -357,7 +488,7 @@ def enclosing(cpu: dict) -> set[str]: else: bucket = "moe" out[bucket][phase][module] = entry - return out, optimizer_us, dpcomm_us + return out, optimizer_us, dpcomm_us, dropped_stall_us def _lists(bd: dict) -> dict: @@ -380,10 +511,12 @@ def _fwd_total(buckets: dict) -> float: 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 = parse_trace(path, ga) + 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 @@ -408,13 +541,24 @@ def build(model: str, traces: dict[str, Path], ga: int = 2) -> dict[str, Any]: # 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 (B=1) Megatron-convention FLOPs at capture seq; site multiplies by cr_layer_counts x GA x DP", + "note": "per-layer and MTP (B=1) Megatron-convention FLOPs at capture seq; site multiplies by GA x DP", } return { @@ -425,10 +569,13 @@ def build(model: str, traces: dict[str, Path], ga: int = 2) -> dict[str, Any]: "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." + "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": { @@ -441,12 +588,13 @@ def build(model: str, traces: dict[str, Path], ga: int = 2) -> dict[str, Any]: "optimizer": "adam", "distributed_optimizer": False, "recompute": "off", - "measured_iter_time_ms": None, + "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"]), + "total_params": model_total_params(model, cfg["num_layers"], mtp_num_layers), }, "hardware": DEFAULT_HARDWARE, "layers": layers, @@ -455,9 +603,9 @@ def build(model: str, traces: dict[str, Path], ga: int = 2) -> dict[str, Any]: "type": "adam", "measured_params": None, "time_us": optimizer_us, - "bytes_per_param": 18, + "bytes_per_param": 30, "class": "memory_bound", - "note": "measured one-layer optimizer-step kernel time on this rank; the site scales by per-rank params", + "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, diff --git a/deepseek-v4/projection/tools/v4_flops.py b/deepseek-v4/projection/tools/v4_flops.py index f9739b5fd..34329866e 100644 --- a/deepseek-v4/projection/tools/v4_flops.py +++ b/deepseek-v4/projection/tools/v4_flops.py @@ -15,6 +15,7 @@ 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, @@ -126,6 +127,16 @@ def _hc_mixer(s, p): 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] @@ -140,16 +151,41 @@ def layer_fmac(model: str, cr: int, seq: int) -> dict[str, float]: } -def nonlayer_fmac(model: str, seq: int) -> dict[str, float]: +def nonlayer_fmac(model: str, seq: int, mtp_num_layers: int = 0) -> dict[str, float]: p = MODEL_PARAMS[model] - return {"logits": seq * p["hidden"] * p["vocab"]} + 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) -> int: +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.""" + 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 = ( @@ -164,7 +200,11 @@ def model_total_params(model: str, num_layers: int) -> int: + SWIGLU * p["hidden"] * p["shared_ffn"] + p["hidden"] * p["experts"] ) - return int(num_layers * (attn + moe) + 2 * p["vocab"] * p["hidden"]) + 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). @@ -184,6 +224,11 @@ 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] diff --git a/examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml index 5861761fd..0731852f0 100644 --- a/examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml @@ -98,40 +98,39 @@ modules: # 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:false} + 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_mlp: false - use_turbo_rms_norm: false - use_turbo_deepep: ${PRIMUS_USE_TURBO_DEEPEP:false} + use_turbo_grouped_mlp: true + use_turbo_rms_norm: true + use_turbo_deepep: ${PRIMUS_USE_TURBO_DEEPEP:true} moe_shared_expert_overlap: false moe_router_dtype: fp32 - # plan-4 P25 / P26: in-tree Primus Triton kernels for V4 attention. - # Precedence in DeepseekV4Attention.forward: - # use_turbo_attention > use_v4_tilelang_attention > use_v4_triton_attention > eager (cr ∈ {0, 128}) - # use_v4_tilelang_csa_attention > use_v4_triton_csa_attention > eager (cr == 4) - use_v4_triton_attention: ${PRIMUS_USE_V4_TRITON_ATTENTION:false} - use_v4_triton_csa_attention: ${PRIMUS_USE_V4_TRITON_CSA_ATTENTION:false} + # 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-8 tilelang attention kernels (OPTIONAL — tilelang is NOT - # bundled in the default Primus container). When enabled, the - # V4 attention dispatcher routes through the tilelang FWD/BWD; - # when tilelang is not installed, the dispatcher prints a single - # rank-0 warning and falls back to Triton -- training continues - # without error. Replaces the legacy PRIMUS_V4_TILELANG_ATTN env - # knob (no longer consulted). - use_v4_tilelang_attention: ${PRIMUS_USE_V4_TILELANG_ATTENTION:false} - use_v4_tilelang_csa_attention: ${PRIMUS_USE_V4_TILELANG_CSA_ATTENTION: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: false + moe_use_fused_router_with_aux_score: true moe_permute_fusion: true # Cross entropy flags 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..eec5b4283 --- /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_mlp: 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_parallel_linear: true + + moe_router_padding_for_quantization: true diff --git a/examples/run_pretrain.sh b/examples/run_pretrain.sh index 4490cf368..c05f3763b 100755 --- a/examples/run_pretrain.sh +++ b/examples/run_pretrain.sh @@ -292,8 +292,11 @@ LOG_INFO "" # ----------------- AMD-specific GPU optimizations ----------------- -# Enable system DMA engine (SDMA) on AMD GPUs for better IO throughput -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 diff --git a/phase5_fp4_indexer_probe.py b/phase5_fp4_indexer_probe.py new file mode 100644 index 000000000..12cda4400 --- /dev/null +++ b/phase5_fp4_indexer_probe.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python +"""Phase 5 probe: FP4 CSA-indexer QK selection-quality + grad-flow check. + +Runs the V4 Indexer in BF16 vs MXFP4-QK (PRIMUS_INDEXER_FP4) on identical +weights/inputs and reports: + * top-k selection OVERLAP (paper's quality metric: do we pick the same + compressed positions?), per query, averaged. + * score cosine (numerical drift of I_{t,s}). + * STE backward sanity (grads finite, non-zero through the FP4 quant). + +Usage: + PRIMUS_INDEXER_FP4=0 not needed — the probe toggles the env itself. + python phase5_fp4_indexer_probe.py +""" +import os + +import torch + +# Import after we can toggle env per-call. +from primus.backends.megatron.core.transformer import indexer as idx_mod +from primus.backends.megatron.core.transformer.indexer import Indexer + +torch.manual_seed(0) + +DEV = "cuda" +DT = torch.bfloat16 + +B, S, D = 2, 2048, 2048 +COMPRESS_RATIO = 4 +INDEX_HEAD_DIM = 128 +INDEX_N_HEADS = 64 +INDEX_TOPK = 64 + + +def build(): + torch.manual_seed(0) + m = Indexer( + hidden_size=D, + index_head_dim=INDEX_HEAD_DIM, + index_n_heads=INDEX_N_HEADS, + index_topk=INDEX_TOPK, + compress_ratio=COMPRESS_RATIO, + ).to(DEV, DT) + return m + + +def run(m, hidden, fp4: bool): + os.environ["PRIMUS_INDEXER_FP4"] = "1" if fp4 else "0" + assert idx_mod._indexer_fp4_enabled() is fp4 + h = hidden.clone().requires_grad_(True) + topk_idxs, topk_scores = m(h) + return topk_idxs, topk_scores, h + + +def main(): + assert torch.cuda.is_available(), "needs MI355X (gfx950) for MXFP4" + m = build() + hidden = torch.randn(B, S, D, device=DEV, dtype=DT) * 0.5 + + idx_bf16, sc_bf16, _ = run(m, hidden, fp4=False) + idx_fp4, sc_fp4, h_fp4 = run(m, hidden, fp4=True) + + # ---- selection overlap (ignore -1 padding/masked slots) ---- + valid = idx_bf16 >= 0 # [B,S,K] + overlaps = [] + Bn, Sn, K = idx_bf16.shape + a = idx_bf16.reshape(-1, K) + b = idx_fp4.reshape(-1, K) + v = valid.reshape(-1, K) + for i in range(a.shape[0]): + sa = set(a[i][v[i]].tolist()) + if not sa: + continue + sb = set(b[i][b[i] >= 0].tolist()) + overlaps.append(len(sa & sb) / len(sa)) + overlap = sum(overlaps) / max(len(overlaps), 1) + + # ---- score cosine on valid (finite) entries ---- + fa = sc_bf16.reshape(-1).float() + fb = sc_fp4.reshape(-1).float() + fin = torch.isfinite(fa) & torch.isfinite(fb) + cos = torch.nn.functional.cosine_similarity(fa[fin], fb[fin], dim=0).item() + + # ---- STE backward sanity through the FP4 quant ---- + loss = sc_fp4[torch.isfinite(sc_fp4)].sum() + loss.backward() + g = h_fp4.grad + grad_ok = bool(torch.isfinite(g).all()) and g.abs().sum().item() > 0 + + print(f"[phase5] shapes: idx {tuple(idx_bf16.shape)} P={S // COMPRESS_RATIO} topk={INDEX_TOPK}") + print(f"[phase5] top-k selection overlap (FP4 vs BF16): {overlap:.4f}") + print(f"[phase5] score cosine (FP4 vs BF16): {cos:.4f}") + print(f"[phase5] STE grad through FP4 QK finite+nonzero: {grad_ok}") + # Selection quality is the gate; the indexer only needs the same top-k. + print(f"[phase5] VERDICT: {'PASS' if overlap > 0.9 and grad_ok else 'CHECK'}") + + +if __name__ == "__main__": + main() diff --git a/primus/backends/megatron/core/extensions/primus_turbo.py b/primus/backends/megatron/core/extensions/primus_turbo.py index 167b647e3..63a907d3e 100644 --- a/primus/backends/megatron/core/extensions/primus_turbo.py +++ b/primus/backends/megatron/core/extensions/primus_turbo.py @@ -3,37 +3,68 @@ # # See LICENSE for license information. ############################################################################### -from contextlib import contextmanager -from typing import Callable, List, Optional, Tuple +import contextlib +import gc +import os +from contextlib import contextmanager, nullcontext +from typing import Callable, Iterable, List, Optional, Tuple, Union -import primus_turbo.pytorch as pt +import primus_turbo.pytorch as primus_turbo_torch import torch import torch.distributed as dist -import torch.nn.functional as F import transformer_engine as te -from megatron.core import tensor_parallel -from megatron.core.extensions.transformer_engine import TELinear, condition_init_method +from megatron.core.enums import Fp4Recipe, Fp8Recipe +from megatron.core.extensions.transformer_engine import ( + TEColumnParallelLinear, + TEGroupedLinear, + TELayerNormColumnParallelLinear, + TELinear, + TEQuantizationParams, + TEQuantizationRecipe, + TERowParallelLinear, + condition_init_method, +) from megatron.core.model_parallel_config import ModelParallelConfig from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.parallel_state import ( get_context_parallel_group, get_hierarchical_context_parallel_groups, get_tensor_model_parallel_group, - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, ) from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.tensor_parallel.layers import ( - ColumnParallelLinear, - _initialize_affine_weight_cpu, -) from megatron.core.transformer.enums import AttnMaskType -from megatron.core.transformer.moe.experts import TEGroupedMLP from megatron.core.transformer.moe.token_dispatcher import MoETokenDispatcher from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint -from megatron.core.utils import divide, get_tensor_model_parallel_group_if_none +from megatron.core.utils import get_pg_size from megatron.training.global_vars import get_args + +# QuantizedTensor / QuantizedTensorPair are only used in the FP8/FP4 weight +# quantization paths (added in PR #735). Older primus_turbo 0.2.0 builds shipped +# in the rocm/primus v26.2 / v26.3 containers do not export them yet. Keep the +# module importable so the BF16 turbo attention / linear paths still work, and only +# fail (with a clear AttributeError on None) if an FP8 quantization path is hit. +try: + from primus_turbo.pytorch.core import QuantizedTensor as PrimusTurboQuantizedTensor + from primus_turbo.pytorch.core import ( + QuantizedTensorPair as PrimusTurboQuantizedTensorPair, + ) +except (ImportError, ModuleNotFoundError): + PrimusTurboQuantizedTensor = None + PrimusTurboQuantizedTensorPair = None + +# ScalingRecipe was renamed to MXScalingRecipe in primus_turbo 0.2.0; keep a fallback +# alias so the module imports against both old and new builds. +try: + from primus_turbo.pytorch.core.low_precision import ScalingRecipe +except (ImportError, ModuleNotFoundError): + from primus_turbo.pytorch.core.low_precision import MXScalingRecipe as ScalingRecipe + +try: + from primus_turbo.pytorch.core.quantized_tensor import create_quantized_weight +except (ImportError, ModuleNotFoundError): + create_quantized_weight = None + from primus_turbo.pytorch.core.low_precision import ( Float4QuantConfig, Float8QuantConfig, @@ -44,16 +75,260 @@ check_fp8_support, check_mxfp4_support, check_mxfp8_support, + float4_e2m1fn_x2, + float8_e4m3, ) from torch import Tensor from transformer_engine.pytorch.constants import dist_group_type from transformer_engine.pytorch.fp8 import DelayedScaling, FP8GlobalStateManager, Recipe -from primus.backends.megatron.core.extensions._triton.stack_grouped_weight import ( - stack_grouped_weight, -) from primus.core.pipeline_parallel.handler.offload_handler import OFFLOAD_BUFFER +try: + import triton + import triton.language as tl + + _HAVE_TRITON = True +except (ImportError, ModuleNotFoundError): + _HAVE_TRITON = False + +_dummy_wgrads = {} + + +if _HAVE_TRITON: + + @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 _triton_inplace_add_(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 Triton is unavailable or the layout is + unsupported (non-contiguous / shape mismatch). The write is in-place on + ``dst``'s storage, so ``dst`` must be contiguous. + """ + if not _HAVE_TRITON or 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 + + +def _get_dummy_wgrad(shape: list, dtype: torch.dtype, zero=False) -> torch.Tensor: + """Returns a dummy tensor of given shape. + + Supports arbitrary rank (2D for plain Linear weights, 3D for stacked + grouped-linear weights ``(num_gemms, out_features, in_features)``, etc.). + Tensors are cached by ``(shape, dtype)`` so each distinct weight layout + only allocates one persistent buffer that gets reused across steps. + """ + global _dummy_wgrads + key = (tuple(shape), dtype) + if key not in _dummy_wgrads: + _dummy_wgrads[key] = torch.empty( + shape, + dtype=dtype, + device="cuda", + requires_grad=False, + ) + if zero: + _dummy_wgrads[key].fill_(0) + return _dummy_wgrads[key].detach() + + +class _MainGradShim: + """Per-expert handle for primus_turbo's ``fused_grouped_wgrad`` over a + *consolidated* ``[E, N, K]`` grouped-expert weight (OPT-1). + + ``PrimusTurboGroupedLinear`` keeps one consolidated weight with a single + ``main_grad`` block, but ``fused_grouped_wgrad`` / ``_expert_main_grad_view`` + expect a list of per-expert handles, each exposing a 2-D ``main_grad`` and a + ``grad_added_to_main_grad`` flag. These shims point at the contiguous 2-D + slices ``main_grad[i]`` of the consolidated block, so the grouped GEMM + backward accumulates each expert's wgrad straight into the right slice. + """ + + __slots__ = ("main_grad", "grad_added_to_main_grad") + + def __init__(self, main_grad_slice: torch.Tensor) -> None: + self.main_grad = main_grad_slice + self.grad_added_to_main_grad = False + + +def _bridge_weight_grad( + x: torch.Tensor, weight: torch.nn.Parameter, weight_buffer: PrimusTurboQuantizedTensorPair +): + """Bridge quantized weight gradient to the original weight's ``main_grad``. + + Must be called **before** the gemm so that in the backward pass the gemm + backward fires first (producing the real weight gradient) and then + ``_WeightGradBridge.backward`` receives it, writes it into + ``weight.main_grad``, and emits a dummy wgrad so that ``weight``'s + AccumulateGrad / DDP ``register_grad_ready`` hook fires in the correct + order. + + """ + + class _WeightGradBridge(torch.autograd.Function): + + @staticmethod + def forward(ctx, x, weight, quantized_weight, quantized_weight_trans): + ctx.save_for_backward(weight) + return x, quantized_weight, quantized_weight_trans + + @staticmethod + def backward(ctx, grad_x, grad_quantized_weight, grad_quantized_weight_trans): + (weight,) = ctx.saved_tensors + assert hasattr(weight, "main_grad"), "weight.main_grad should be set before backward pass." + assert hasattr( + weight, "grad_added_to_main_grad" + ), "weight.grad_added_to_main_grad don't have grad_added_to_main_grad attribute." + + # NOTE: Set weight.grad_added_to_main_grad to True to avoid adding the + # quantized weight gradient to main_grad twice. + if grad_quantized_weight is None: + # OPT-1 fused path: the grouped GEMM backward already accumulated the + # expert wgrad straight into main_grad (under fused_grouped_wgrad) and + # returned grad_b=None, so there is nothing to add here -- just flag it. + weight.grad_added_to_main_grad = True + else: + # `is_gfx1250` only exists in newer primus_turbo builds (it gates a + # gfx1250-specific elementwise-add workaround). Older / feature-branch + # primus_turbo that predate it (e.g. the flydsl sparse-MLA attention branch) + # don't define it; treat a missing symbol as False so those builds still work + # on non-gfx1250 archs (gfx942 / gfx950) instead of raising ImportError. + try: + from primus_turbo.pytorch.core.utils import is_gfx1250 + + _use_triton_inplace_add = is_gfx1250() + except ImportError: + _use_triton_inplace_add = False + + if _use_triton_inplace_add: + # NOTE: The bandwith of torch's elementwise add kernel has issue. Use triton to temporary workaround for gfx1250. + _triton_inplace_add_(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 + + assert isinstance( + weight_buffer, PrimusTurboQuantizedTensorPair + ), "weight_buffer must be a PrimusTurboQuantizedTensorPair" + assert weight_buffer.data is not None, "weight_buffer.data must not be None" + + x, quantized_weight, quantized_weight_trans = _WeightGradBridge.apply( + x, weight, weight_buffer.data, weight_buffer.data_t + ) + + # wrapper quantized_weight and quantized_weight_trans into PrimusTurboQuantizedTensorPair + return x, PrimusTurboQuantizedTensorPair(data=quantized_weight, data_t=quantized_weight_trans) + + +def _maybe_create_quantized_weight_buffers( + weight: torch.Tensor, + dest_dtype: torch.dtype, + quant_config: "PrimusTurboQuantConfig", + disable_parameter_transpose_cache: bool, +): + """Quantize ``weight`` into a rowwise buffer plus an optional transposed + (colwise) buffer, returning ``(rowwise, colwise_or_None)``. + + Prefers primus_turbo's ``create_quantized_weight`` helper, which picks the + scaling recipe from the quant config and handles per-granularity transpose. + Falls back to a manual rowwise/colwise quantize on older primus_turbo builds + that do not export ``create_quantized_weight`` yet. + """ + quant_config_internal = quant_config.data() + need_weight_transpose_cache = not disable_parameter_transpose_cache + + if create_quantized_weight is not None: + return create_quantized_weight( + weight, + dest_dtype, + quant_config_internal, + need_weight_transpose_cache=need_weight_transpose_cache, + ) + + # TODO(ruibin): Remove this fallback path once create_quantized_weight is + # always available in the shipped primus_turbo build. + def _weight_scaling_recipe(quant_config: Union[Float4QuantConfig, Float8QuantConfig]) -> ScalingRecipe: + if isinstance(quant_config, Float4QuantConfig): + weight_scaling_recipe = ScalingRecipe( + use_2d_block=True, + shuffle_scale=quant_config.use_preshuffle, + shuffle_out=quant_config.use_preshuffle, + ) + + if isinstance(quant_config, Float8QuantConfig): + if quant_config.granularity in [ScalingGranularity.BLOCKWISE, ScalingGranularity.MX_BLOCKWISE]: + weight_scaling_recipe = ScalingRecipe(use_2d_block=True) + else: + weight_scaling_recipe = ScalingRecipe() + + return weight_scaling_recipe + + quantized_weight_rowwise = PrimusTurboQuantizedTensor.quantize( + weight, + dest_dtype=dest_dtype, + granularity=quant_config.granularity, + block_size=quant_config.block_size, + scaling_recipe=_weight_scaling_recipe(quant_config), + axis=-1, + ) + + quantized_weight_colwise = None + if need_weight_transpose_cache: + granularity = quant_config.granularity + if granularity == ScalingGranularity.TENSORWISE: + quantized_weight_colwise = quantized_weight_rowwise.transpose(-2, -1) + elif granularity == ScalingGranularity.ROWWISE: + # NOTE: rowwise quantization not support transpose, so we need to quantize the transposed weight manually. + quantized_weight_colwise = PrimusTurboQuantizedTensor.quantize( + weight.transpose(-2, -1), + dest_dtype=dest_dtype, + granularity=quant_config.granularity, + block_size=quant_config.block_size, + scaling_recipe=_weight_scaling_recipe(quant_config), + axis=-2, + ) + elif granularity in [ScalingGranularity.BLOCKWISE, ScalingGranularity.MX_BLOCKWISE]: + quantized_weight_colwise = PrimusTurboQuantizedTensor.quantize( + weight, + dest_dtype=dest_dtype, + granularity=quant_config.granularity, + block_size=quant_config.block_size, + scaling_recipe=_weight_scaling_recipe(quant_config), + # axis=-2 means quant weight along axis 2 which will get a transposed quantized weight. + axis=-2, + ) + else: + raise ValueError(f"Unsupported granularity: {granularity}") + + return quantized_weight_rowwise, quantized_weight_colwise + def _call_fp8_autocast_enter( *, @@ -88,8 +363,17 @@ def _call_fp8_autocast_exit(enabled: bool, *, _graph: bool) -> None: exit_fn(enabled, _graph=_graph) -def use_split_wgrad_op(): +def _is_fp4_or_fp8_enabled(): + return ( + PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled() + or PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp4_enabled() + ) + + +def _use_split_wgrad_op(): args = get_args() + + enable_split_wgrad_op = False if args.patch_primus_pipeline and args.pp_algorithm in [ "zero-bubble", "zero-bubble-heuristic", @@ -97,10 +381,17 @@ def use_split_wgrad_op(): "v-half", "v-min", ]: - return True + enable_split_wgrad_op = True + elif args.patch_zero_bubble and args.enable_zero_bubble: - return True - return False + enable_split_wgrad_op = True + + if enable_split_wgrad_op: + assert ( + not _is_fp4_or_fp8_enabled() + ), "split wgrad op is not supported when turbo fp8 or fp4 is enabled." + + return enable_split_wgrad_op class PrimusTurboQuantConfig: @@ -332,6 +623,91 @@ def primus_turbo_fp4_autocast( _call_fp8_autocast_exit(enabled, _graph=_graph) +def _get_fp8_autocast_for_quant_recipe(qrecipe: TEQuantizationRecipe): + if FP8GlobalStateManager.is_fp8_enabled(): + if not qrecipe.override_quantized_autocast: + return nullcontext() + else: + if not qrecipe.override_nonquantized_autocast: + return nullcontext() + + if qrecipe.fp8_quantization_recipe is None and qrecipe.fp4_quantization_recipe is None: + # Force BF16 for this layer and override autocast + return primus_turbo_fp8_autocast(enabled=False, enabled_turbo=False) + else: + if ( + qrecipe.fp8_quantization_recipe == Fp8Recipe.custom + or qrecipe.fp4_quantization_recipe == Fp4Recipe.custom + ): + assert qrecipe.custom_recipe_factory is not None + assert False, "Custom recipe is not supported for Primus-Turbo" + + elif qrecipe.fp8_quantization_recipe is not None: + from primus.backends.megatron.core.fp8_utils import ( + MXFP8_SCALING_BLOCK_SIZE, + SCALING_BLOCK_SIZE, + ) + + if qrecipe.fp8_format == "e4m3": + fp8_format = Format.E4M3 + elif qrecipe.fp8_format == "hybrid": + fp8_format = Format.HYBRID + else: + raise ValueError(f"Unhandled fp8_format {qrecipe.fp8_format}") + + if qrecipe.fp8_quantization_recipe == Fp8Recipe.tensorwise: + quant_recipe = PrimusTurboQuantConfig( + granularity=ScalingGranularity.TENSORWISE, format=fp8_format + ) + elif qrecipe.fp8_quantization_recipe == Fp8Recipe.blockwise: + quant_recipe = PrimusTurboQuantConfig( + granularity=ScalingGranularity.BLOCKWISE, format=fp8_format, block_size=SCALING_BLOCK_SIZE + ) + elif qrecipe.fp8_quantization_recipe == Fp8Recipe.mxfp8: + quant_recipe = PrimusTurboQuantConfig( + granularity=ScalingGranularity.MX_BLOCKWISE, + format=fp8_format, + block_size=MXFP8_SCALING_BLOCK_SIZE, + scale_dtype=ScaleDtype.E8M0, + ) + else: + raise ValueError(f"Unhandled fp8 recipe: {qrecipe.fp8_quantization_recipe}") + + return primus_turbo_fp8_autocast( + enabled=False, enabled_turbo=True, turbo_quant_config=quant_recipe + ) + else: + # Fp4 configured. + if qrecipe.fp4_quantization_recipe == Fp4Recipe.nvfp4: + assert False, "NVFP4 is not supported for Primus-Turbo" + elif qrecipe.fp4_quantization_recipe == Fp4Recipe.mxfp4: + from primus.backends.megatron.core.fp4_utils import ( + MXFP4_SCALING_BLOCK_SIZE, + ) + + quant_recipe = PrimusTurboQuantConfig( + granularity=ScalingGranularity.MX_BLOCKWISE, + format=Format.E2M1_X2, + block_size=MXFP4_SCALING_BLOCK_SIZE, + scale_dtype=ScaleDtype.E8M0, + ) + else: + raise ValueError(f"Unhandled fp4 recipe: {qrecipe.fp4_quantization_recipe}") + + return primus_turbo_fp4_autocast( + enabled=False, enabled_turbo=True, turbo_quant_config=quant_recipe + ) + + +def _get_fp8_autocast_for_quant_params(qparams: TEQuantizationParams | None, training: bool): + if qparams is None: + return nullcontext() + elif not training and qparams.evaluation_recipe is not None: + return _get_fp8_autocast_for_quant_recipe(qparams.evaluation_recipe) + else: + return _get_fp8_autocast_for_quant_recipe(qparams.training_recipe) + + class PrimusTurboAttention(te.pytorch.DotProductAttention): """ Wrapper for the Transformer-Engine's `DotProductAttention` layer that also @@ -396,6 +772,9 @@ def __init__( ) _use_sink_attention = False + # set deterministic flag + self.deterministic_mode = args.deterministic_mode + # Store for later use after super().__init__() self._init_sink_attention = _use_sink_attention self._num_heads_for_sinks = self.config.num_attention_heads @@ -403,15 +782,15 @@ def __init__( self.offload = args.offload and "attn" in args.offload_ops if args.enable_turbo_attention_float8: self.attn = ( - pt.ops.flash_attn_fp8_usp_func + primus_turbo_torch.ops.flash_attn_fp8_usp_func if self.config.context_parallel_size > 1 - else pt.ops.flash_attn_fp8_func + else primus_turbo_torch.ops.flash_attn_fp8_func ) else: self.attn = ( - pt.ops.flash_attn_usp_func + primus_turbo_torch.ops.flash_attn_usp_func if self.config.context_parallel_size > 1 - else pt.ops.flash_attn_func + else primus_turbo_torch.ops.flash_attn_func ) if pg_collection is None: # For backward compatibility, remove in v0.14 and raise error @@ -488,8 +867,6 @@ def forward( packed_seq_params: PackedSeqParams = None, ): """Forward.""" - SUPPORTED_QKV_FORMATS = "sbhd" - packed_seq_kwargs = ( {key: getattr(packed_seq_params, key) for key in self.kept_packed_seq_params} if packed_seq_params is not None @@ -497,11 +874,6 @@ def forward( ) qkv_format = packed_seq_kwargs.get("qkv_format", self.qkv_format) - assert ( - qkv_format in SUPPORTED_QKV_FORMATS - ), f"qkv_format only support {SUPPORTED_QKV_FORMATS}, but got {qkv_format}" - # NOTE(ruibin): The layout of q, k and v is (S, B, H, D). But attn accept the shape of qkv is (B, S, H, D). - query, key, value = [x.permute(1, 0, 2, 3) for x in (query, key, value)] mask_type = attn_mask_type.name if mask_type == AttnMaskType.causal.name: causal = True @@ -523,7 +895,9 @@ def forward( sink_tensor = None window_size = (-1, -1) - if self.use_sink_attention and self.sinks is not None: + use_sink_attn = self.use_sink_attention and self.sinks is not None + + if use_sink_attn: sink_tensor = self.sinks # Apply sliding window based on layer pattern (gpt-oss: even layers only) @@ -535,11 +909,28 @@ def forward( window_size = (self.sink_sliding_window, 0) else: window_size = (self.sink_sliding_window, 0) + if self.offload: OFFLOAD_BUFFER.add_offload_tensor(f"attn_q", query) OFFLOAD_BUFFER.add_offload_tensor(f"attn_k", key) OFFLOAD_BUFFER.add_offload_tensor(f"attn_v", value) + # NOTE: query, key, value maybe a view of the original tensor, call contiguous to copy a new tensor + # and let torch allocator can release the original tensor. + if torch.is_grad_enabled(): + query = query.contiguous() if query.requires_grad else query + key = key.contiguous() if key.requires_grad else key + value = value.contiguous() if value.requires_grad else value + + if qkv_format == "sbhd": + query = query.permute(1, 0, 2, 3) + key = key.permute(1, 0, 2, 3) + value = value.permute(1, 0, 2, 3) + elif qkv_format == "bhsd": + query = query.permute(0, 2, 1, 3) + key = key.permute(0, 2, 1, 3) + value = value.permute(0, 2, 1, 3) + o = self.attn( query, key, @@ -550,15 +941,18 @@ def forward( window_size=window_size, bias=None, alibi_slopes=None, - deterministic=False, + deterministic=self.deterministic_mode, return_lse=False, return_attn_probs=False, sink=sink_tensor, # PR 208: pass sink tensor to Primus-Turbo **self.attn_kwargs, ) - # NOTE(ruibin): The output of attn is BSHD. Use permute to convert the layout to SBHD. - o = o.permute(1, 0, 2, 3).contiguous() + if qkv_format == "sbhd": + o = o.permute(1, 0, 2, 3) + elif qkv_format == "bhsd": + o = o.permute(0, 2, 1, 3) + o = o.view(o.shape[0], o.shape[1], -1) return o @@ -585,44 +979,31 @@ def __init__( symmetric_ar_type: Optional[str] = None, tp_group: Optional[torch.distributed.ProcessGroup] = None, ): - if parallel_mode != "duplicated": - raise ValueError(f"{__class__.__name__} only support `parallel_mode=duplicated`. ") - - tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert) - - if use_split_wgrad_op(): - from .zbpp_gemm import gemm_with_weight_gradient_store - - self.gemm = lambda a, b, bias=None: gemm_with_weight_gradient_store(a, b, bias=bias) - else: - self.gemm = lambda a, b, trans_a=False, trans_b=True, out_dtype=None: pt.ops.gemm( - a, b, trans_a=trans_a, trans_b=trans_b, out_dtype=out_dtype - ) - args = get_args() - self.offload = args.offload and "basic_gemm" in args.offload_ops + self.offload = args.offload and "parallel_gemm" in args.offload_ops assert not self.offload, "gemm offload still have some problems" super().__init__( input_size=input_size, output_size=output_size, - parallel_mode="row", + parallel_mode=parallel_mode, config=config, - init_method=( - condition_init_method(config, init_method) - if not config.use_cpu_initialization - else lambda w: None - ), + init_method=init_method, bias=bias, skip_bias_add=skip_bias_add, - skip_weight_param_allocation=False, - # We don't currently use this for row parallel layers # pylint: disable=line-too-long - is_expert=is_expert, + skip_weight_param_allocation=skip_weight_param_allocation, tp_comm_buffer_name=tp_comm_buffer_name, - symmetric_ar_type=config.symmetric_ar_type, + is_expert=is_expert, + symmetric_ar_type=symmetric_ar_type, tp_group=tp_group, ) + tp_size = get_pg_size(self._tp_group) + assert tp_size == 1, "PrimusTurboLinear only supports tensor parallel size = 1" + + self.register_buffer("quantized_weight_buffer", None, persistent=False) + self.register_buffer("quantized_weight_t_buffer", None, persistent=False) + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): """Sharding along axis 1, bias not sharded""" state_dict = self.state_dict(prefix="", keep_vars=True) @@ -634,55 +1015,116 @@ def __repr__(self): f"out_features={self.out_features}, bias={self.use_bias}, TP={self.tp_size})" ) - def forward( + def forward(self, x: torch.Tensor): + _is_first_microbatch = self.is_first_microbatch + + # Rewrite quant context + quant_context = _get_fp8_autocast_for_quant_params(self.te_quant_params, self.training) + + with quant_context: + out = self.forward_internal(x, _is_first_microbatch) + + self.is_first_microbatch = False + + return out + + def forward_internal( self, - input_: torch.Tensor, + x: torch.Tensor, + is_first_microbatch: bool = False, ): - # weights = [getattr(self, name) for name in self.weight_names] - # weights = torch.cat(weights, dim=0) # or set weights = self._parameters['weight'] - weights = self._parameters["weight"] + weight = self._parameters["weight"] if self.use_bias: bias_tensor = torch.cat([getattr(self, name) for name in self.bias_names]) - original_shape = input_.size() - if not input_.is_contiguous(): - input_ = input_.contiguous() - input_ = input_.view(-1, original_shape[-1]) + original_shape = x.size() + if not x.is_contiguous(): + x = x.contiguous() + x = x.view(-1, original_shape[-1]) if self.offload: - OFFLOAD_BUFFER.add_offload_tensor(f"linear_input", input_) + OFFLOAD_BUFFER.add_offload_tensor(f"linear_input", x) - if PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled(): - quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() - if quant_config.current_scaling() or quant_config.block_scaling() or quant_config.mxfp8_scaling(): - fp8_gemm = pt.ops.gemm_fp8 - else: - raise ValueError("Not support quant config.") - - out = fp8_gemm( - input_, weights, 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() - if quant_config.mxfp4_scaling(): - fp4_gemm = pt.ops.gemm_fp4 - else: - raise ValueError("Not support quant config.") + if _use_split_wgrad_op(): + from .zbpp_gemm import gemm_with_weight_gradient_store - out = fp4_gemm( - input_, weights, trans_a=False, trans_b=True, out_dtype=None, config=quant_config.data() - ) + out = gemm_with_weight_gradient_store(x, weight, bias=None) else: - out = self.gemm(input_, weights) + if PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled(): + quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() + assert ( + 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( + weight, + float8_e4m3, + 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_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( + 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(), + ) + else: + out = primus_turbo_torch.ops.gemm(x, weight, trans_a=False, trans_b=True, out_dtype=None) out = out.view(original_shape[0], original_shape[1], -1) - if self.te_return_bias: - return out, bias_tensor + if self.use_bias: - return out + bias_tensor, None + out = out + bias_tensor + return out, None -class PrimusTurboRowParallelLinear(TELinear): +class PrimusTurboRowParallelLinear(TERowParallelLinear): """ Wrapper for the Transformer-Engine's `Linear` layer but specialized similar to megatron's `RowParallelLinear` layer. @@ -709,67 +1151,24 @@ def __init__( self.offload = args.offload and "row_parallel_gemm" in args.offload_ops assert not self.offload, "gemm offload still have some problems" - tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert) - - if use_split_wgrad_op(): - from .zbpp_gemm import gemm_with_weight_gradient_store - - self.gemm = lambda a, b, bias=None: gemm_with_weight_gradient_store(a, b, bias=bias) - else: - self.gemm = lambda a, b, trans_a=False, trans_b=True, out_dtype=None: pt.ops.gemm( - a, b, trans_a=trans_a, trans_b=trans_b, out_dtype=out_dtype - ) - super().__init__( input_size=input_size, output_size=output_size, - parallel_mode="row", config=config, - init_method=( - condition_init_method(config, init_method) - if not config.use_cpu_initialization - else lambda w: None - ), + init_method=init_method, bias=bias, + input_is_parallel=input_is_parallel, skip_bias_add=skip_bias_add, - skip_weight_param_allocation=False, - # We don't currently use this for row parallel layers # pylint: disable=line-too-long is_expert=is_expert, tp_comm_buffer_name=tp_comm_buffer_name, - symmetric_ar_type=config.symmetric_ar_type, tp_group=tp_group, ) - # Manual weight initialization when use_cpu_initialization=True - if config.use_cpu_initialization: - world_size = get_tensor_model_parallel_world_size() - rank = get_tensor_model_parallel_rank() - input_size_per_partition = divide(input_size, world_size) - - self.master_weight = _initialize_affine_weight_cpu( - self.weight, - output_size, - input_size, - input_size_per_partition, - 1, # partition_dim (row parallel partitions along input dimension) - init_method=condition_init_method(config, init_method), - stride=1, - return_master_weight=False, - params_dtype=config.params_dtype, - rank=rank, - world_size=world_size, - skip_set_tensor_parallel_attributes=True, - ) + tp_size = get_pg_size(self._tp_group) + assert tp_size == 1, "PrimusTurboRowParallelLinear only supports tensor parallel size = 1" - # Bias initialization - if bias: - with torch.no_grad(): - bias_tensor = torch.cat([getattr(self, name) for name in self.bias_names]) - bias_tensor.zero_() - # Set allreduce attribute for distributed training - for bias_name in self.bias_names: - bias_param = getattr(self, bias_name) - setattr(bias_param, "allreduce", True) + self.register_buffer("quantized_weight_buffer", None, persistent=False) + self.register_buffer("quantized_weight_t_buffer", None, persistent=False) def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): """Sharding along axis 1, bias not sharded""" @@ -782,59 +1181,121 @@ def __repr__(self): f"out_features={self.out_features}, bias={self.use_bias}, TP={self.tp_size})" ) - def forward( + def forward(self, x: torch.Tensor): + _is_first_microbatch = self.is_first_microbatch + + # Rewrite quant context + quant_context = _get_fp8_autocast_for_quant_params(self.te_quant_params, self.training) + + with quant_context: + out = self.forward_internal(x, _is_first_microbatch) + + self.is_first_microbatch = False + + return out + + def forward_internal( self, - input_: torch.Tensor, + x: torch.Tensor, + is_first_microbatch: bool = False, ): - # weights = [getattr(self, name) for name in self.weight_names] - # weights = torch.cat(weights, dim=0) # or set weights = self._parameters['weight'] - weights = self._parameters["weight"] + weight = self._parameters["weight"] + if self.use_bias: bias_tensor = torch.cat([getattr(self, name) for name in self.bias_names]) - original_shape = input_.size() - if not input_.is_contiguous(): - input_ = input_.contiguous() + original_shape = x.size() + if not x.is_contiguous(): + x = x.contiguous() if self.offload: - OFFLOAD_BUFFER.add_offload_tensor(f"row_parallel_linear_input", input_) + OFFLOAD_BUFFER.add_offload_tensor(f"row_parallel_linear_input", x) - input_ = input_.view(-1, original_shape[-1]) + x = x.view(-1, original_shape[-1]) if self.offload: - OFFLOAD_BUFFER.add_offload_tensor(f"row_parallel_linear_input", input_) + OFFLOAD_BUFFER.add_offload_tensor(f"row_parallel_linear_input", x) - if PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled(): - quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() - if quant_config.current_scaling() or quant_config.block_scaling() or quant_config.mxfp8_scaling(): - fp8_gemm = pt.ops.gemm_fp8 - else: - raise ValueError("Not support quant config.") - - out = fp8_gemm( - input_, weights, 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() - if quant_config.mxfp4_scaling(): - fp4_gemm = pt.ops.gemm_fp4 - else: - raise ValueError("Not support quant config.") + if _use_split_wgrad_op(): + from .zbpp_gemm import gemm_with_weight_gradient_store - out = fp4_gemm( - input_, weights, trans_a=False, trans_b=True, out_dtype=None, config=quant_config.data() - ) + out = gemm_with_weight_gradient_store(x, weight, bias=None) else: - out = self.gemm(input_, weights) + if PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled(): + quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() + assert ( + 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( + weight, + float8_e4m3, + 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_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( + 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(), + ) + else: + out = primus_turbo_torch.ops.gemm(x, weight, trans_a=False, trans_b=True, out_dtype=None) out = out.view(original_shape[0], original_shape[1], -1) - if self.te_return_bias: - return out, bias_tensor + if self.use_bias: - return out + bias_tensor, None + out = out + bias_tensor + return out, None -class PrimusTurboColumnParallelLinear(TELinear): +class PrimusTurboColumnParallelLinear(TEColumnParallelLinear): def __init__( self, input_size: int, @@ -851,72 +1312,30 @@ def __init__( tp_group: Optional[torch.distributed.ProcessGroup] = None, stride: int = 1, # TODO(ruibin): compatible with Megatron-LM. Not used. ): - if gather_output: - raise ValueError(f"{__class__.__name__} layers do not support gather_output = True") - tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert) - args = get_args() self.offload = args.offload and "column_parallel_gemm" in args.offload_ops assert not self.offload, "gemm offload still have some problems" - if use_split_wgrad_op(): - from .zbpp_gemm import gemm_with_weight_gradient_store - - self.gemm = lambda a, b, bias=None: gemm_with_weight_gradient_store(a, b, bias=bias) - else: - self.gemm = lambda a, b, trans_a=False, trans_b=True, out_dtype=None: pt.ops.gemm( - a, b, trans_a=trans_a, trans_b=trans_b, out_dtype=out_dtype - ) - super().__init__( input_size=input_size, output_size=output_size, - parallel_mode="column", config=config, - init_method=( - condition_init_method(config, init_method) - if not config.use_cpu_initialization - else lambda w: None - ), + 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, - symmetric_ar_type=config.symmetric_ar_type, tp_group=tp_group, + stride=stride, ) - # Manual weight initialization when use_cpu_initialization=True - if config.use_cpu_initialization: - world_size = get_tensor_model_parallel_world_size() - rank = get_tensor_model_parallel_rank() - output_size_per_partition = divide(output_size, world_size) - - _ = _initialize_affine_weight_cpu( - self.weight, - output_size, - input_size, - output_size_per_partition, - 0, # partition_dim (column parallel partitions along output dimension) - init_method=condition_init_method(config, init_method), - stride=1, - return_master_weight=False, - params_dtype=config.params_dtype, - rank=rank, - world_size=world_size, - skip_set_tensor_parallel_attributes=True, - ) + tp_size = get_pg_size(self._tp_group) + assert tp_size == 1, "PrimusTurboColumnParallelLinear only supports tensor parallel size = 1" - # Bias initialization - if bias: - with torch.no_grad(): - bias_tensor = torch.cat([getattr(self, name) for name in self.bias_names]) - bias_tensor.zero_() - # Set allreduce attribute for distributed training - for bias_name in self.bias_names: - bias_param = getattr(self, bias_name) - setattr(bias_param, "allreduce", True) + self.register_buffer("quantized_weight_buffer", None, persistent=False) + self.register_buffer("quantized_weight_t_buffer", None, persistent=False) def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): """Sharding along axis 0, bias sharded""" @@ -931,158 +1350,116 @@ def __repr__(self): f"out_features={self.out_features}, bias={self.use_bias}, TP={self.tp_size})" ) - def forward( - self, - input_: torch.Tensor, - ): - # weights = [getattr(self, name) for name in self.weight_names] - # weights = torch.cat(weights, dim=0) # or set weights = self._parameters['weight'] - weights = self._parameters["weight"] - if self.use_bias: - bias_tensor = torch.cat([getattr(self, name) for name in self.bias_names]) - original_shape = input_.size() - if not input_.is_contiguous(): - input_ = input_.contiguous() - input_ = input_.view(-1, original_shape[-1]) + def forward(self, x: torch.Tensor): + _is_first_microbatch = self.is_first_microbatch - if self.offload: - OFFLOAD_BUFFER.add_offload_tensor(f"column_parallel_linear_input", input_) + # Rewrite quant context + quant_context = _get_fp8_autocast_for_quant_params(self.te_quant_params, self.training) - if PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled(): - quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() - if quant_config.current_scaling() or quant_config.block_scaling() or quant_config.mxfp8_scaling(): - fp8_gemm = pt.ops.gemm_fp8 - else: - raise ValueError("Not support quant config.") + with quant_context: + out = self.forward_internal(x, _is_first_microbatch) - out = fp8_gemm( - input_, weights, 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() - if quant_config.mxfp4_scaling(): - fp4_gemm = pt.ops.gemm_fp4 - else: - raise ValueError("Not support quant config.") - - out = fp4_gemm( - input_, weights, trans_a=False, trans_b=True, out_dtype=None, config=quant_config.data() - ) - else: - out = self.gemm(input_, weights) + self.is_first_microbatch = False - out = out.view(original_shape[0], original_shape[1], -1) - if self.te_return_bias: - return out, bias_tensor - if self.use_bias: - return out + bias_tensor, None - return out, None + return out - -class PrimusTurboColumnParallelLinearTorch(ColumnParallelLinear): - """ - Wrapper for the Transformer-Engine's `Linear` layer but specialized similar - to megatron's `ColumnParallelLinear` layer. - """ - - def __init__( + def forward_internal( self, - input_size, - output_size, - *, - config: ModelParallelConfig, - init_method: Callable, - bias=True, - gather_output=False, - stride=1, - keep_master_weight_for_test=False, - skip_bias_add=False, - skip_weight_param_allocation: bool = False, - embedding_activation_buffer: Optional[List[torch.Tensor]] = None, - grad_output_buffer: Optional[List[torch.Tensor]] = None, - is_expert: bool = False, - tp_comm_buffer_name: str = None, # Not used - disable_grad_reduce: bool = False, - tp_group: Optional[torch.distributed.ProcessGroup] = None, + x: torch.Tensor, + is_first_microbatch: bool = False, ): - args = get_args() - self.offload = args.offload and "column_parallel_gemm" in args.offload_ops - assert not self.offload, "gemm offload still have some problems" + weight = self._parameters["weight"] + if self.use_bias: + bias_tensor = torch.cat([getattr(self, name) for name in self.bias_names]) + original_shape = x.size() + if not x.is_contiguous(): + x = x.contiguous() + x = x.view(-1, original_shape[-1]) + + if self.offload: + OFFLOAD_BUFFER.add_offload_tensor(f"column_parallel_linear_input", x) - if use_split_wgrad_op(): + if _use_split_wgrad_op(): from .zbpp_gemm import gemm_with_weight_gradient_store - self.gemm = lambda a, b, bias=None: gemm_with_weight_gradient_store(a, b, bias=bias) + out = gemm_with_weight_gradient_store(x, weight, bias=None) else: - self.gemm = lambda a, b, trans_a=False, trans_b=True, out_dtype=None: pt.ops.gemm( - a, b, trans_a=trans_a, trans_b=trans_b, out_dtype=out_dtype - ) - - super().__init__( - input_size, - output_size, - config=config, - init_method=init_method, - bias=bias, - gather_output=gather_output, - stride=stride, - keep_master_weight_for_test=keep_master_weight_for_test, - skip_bias_add=skip_bias_add, - skip_weight_param_allocation=skip_weight_param_allocation, - embedding_activation_buffer=embedding_activation_buffer, - grad_output_buffer=grad_output_buffer, - is_expert=is_expert, - tp_comm_buffer_name=tp_comm_buffer_name, - disable_grad_reduce=disable_grad_reduce, - tp_group=tp_group, - ) - - def forward( - self, - input_: torch.Tensor, - weight: Optional[torch.Tensor] = None, - runtime_gather_output: Optional[bool] = None, - ): - if weight is None: - weight = self.weight - bias_tensor = self.bias if not self.skip_bias_add else None - - original_shape = input_.size() - if not input_.is_contiguous(): - input_ = input_.contiguous() - input_ = input_.view(-1, original_shape[-1]) - - if self.offload: - OFFLOAD_BUFFER.add_offload_tensor(f"column_parallel_linear_torch_input", input_) + if PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled(): + quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() + assert ( + 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( + weight, + float8_e4m3, + quant_config, + disable_parameter_transpose_cache=self.disable_parameter_transpose_cache, + ) - if PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled(): - quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() - if quant_config.current_scaling() or quant_config.block_scaling() or quant_config.mxfp8_scaling(): - fp8_gemm = pt.ops.gemm_fp8 - else: - raise ValueError("Not support quant config.") + 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( + weight, + float4_e2m1fn_x2, + quant_config, + disable_parameter_transpose_cache=self.disable_parameter_transpose_cache, + ) - out = fp8_gemm( - input_, 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() - if quant_config.mxfp4_scaling(): - fp4_gemm = pt.ops.gemm_fp4 + 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: - raise ValueError("Not support quant config.") + out = primus_turbo_torch.ops.gemm(x, weight, trans_a=False, trans_b=True, out_dtype=None) - out = fp4_gemm( - input_, weight, trans_a=False, trans_b=True, out_dtype=None, config=quant_config.data() - ) - else: - out = self.gemm(input_, weight) out = out.view(original_shape[0], original_shape[1], -1) - return out, bias_tensor + if self.use_bias: + out = out + bias_tensor + + return out, None -class PrimusTurboLayerNormColumnParallelLinear(te.pytorch.LayerNormLinear): +class PrimusTurboLayerNormColumnParallelLinear(TELayerNormColumnParallelLinear): """ Wrapper for the Transformer-Engine's `LayerNormLinear` layer that combines layernorm and linear layers @@ -1102,94 +1479,32 @@ def __init__( skip_weight_param_allocation: bool = False, tp_comm_buffer_name: Optional[str] = None, tp_group: Optional[torch.distributed.ProcessGroup] = None, - stride: int = 1, # TODO(ruibin): compatible with Megatron-LM. Not used. + stride: int = 1, ): args = get_args() - self.config = config self.offload = args.offload and "column_parallel_gemm" in args.offload_ops assert not self.offload, "gemm offload still have some problems" - if gather_output: - raise ValueError("Primus Turbo linear layers do not support gather_output = True") - - if is_expert: - raise ValueError("Primus Turbo linear layers do not yet support MoE") - - if skip_weight_param_allocation: - raise ValueError("Primus Turbo linear layers do not support skip_weight_param_allocation") - - # TODO: For backward compatibility, remove in v0.15. - tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert) - - # TE returns a zero length Tensor when bias=False and - # return_bias=True, but we prefer None. So in that case we - # tell TE to not return the bias, and return None - # ourselves. This way our forward always returns two values - # and we don't have to deal with the zero length Tensor. - self.te_return_bias = skip_bias_add and bias - - if use_split_wgrad_op(): - - from .zbpp_gemm import gemm_with_weight_gradient_store - - self.gemm = lambda a, b, bias=None: gemm_with_weight_gradient_store(a, b, bias=bias) - else: - self.gemm = lambda a, b, trans_a=False, trans_b=True, out_dtype=None: pt.ops.gemm( - a, b, trans_a=trans_a, trans_b=trans_b, out_dtype=out_dtype - ) - super().__init__( - in_features=input_size, - out_features=output_size, - eps=self.config.layernorm_epsilon, - sequence_parallel=self.config.sequence_parallel, - fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion, - tp_group=tp_group if torch.distributed.is_initialized() else None, - tp_size=self.config.tensor_model_parallel_size, - get_rng_state_tracker=None, - init_method=( - condition_init_method(config, init_method) - if not config.use_cpu_initialization - else lambda w: None - ), + input_size, + output_size, + config=config, + init_method=init_method, + gather_output=gather_output, bias=bias, - normalization=self.config.normalization, - return_bias=self.te_return_bias, - parallel_mode="column", - return_layernorm_output=False, - zero_centered_gamma=self.config.layernorm_zero_centered_gamma, + 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, + stride=stride, ) - # Manual weight initialization when use_cpu_initialization=True - if config.use_cpu_initialization: - world_size = get_tensor_model_parallel_world_size() - rank = get_tensor_model_parallel_rank() - output_size_per_partition = divide(output_size, world_size) - - _ = _initialize_affine_weight_cpu( - self.weight, - output_size, - input_size, - output_size_per_partition, - 0, # partition_dim (column parallel partitions along output dimension) - init_method=condition_init_method(config, init_method), - stride=1, - return_master_weight=False, - params_dtype=config.params_dtype, - rank=rank, - world_size=world_size, - skip_set_tensor_parallel_attributes=True, - ) + tp_size = get_pg_size(self._tp_group) + assert tp_size == 1, "PrimusTurboLayerNormColumnParallelLinear only supports tensor parallel size = 1" - # Bias initialization - if bias: - with torch.no_grad(): - bias_tensor = torch.cat([getattr(self, name) for name in self.bias_names]) - bias_tensor.zero_() - # Set allreduce attribute for distributed training - for bias_name in self.bias_names: - bias_param = getattr(self, bias_name) - setattr(bias_param, "allreduce", True) + self.register_buffer("quantized_weight_buffer", None, persistent=False) + self.register_buffer("quantized_weight_t_buffer", None, persistent=False) def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): """Sharding along axis 0, bias sharded""" @@ -1205,233 +1520,447 @@ def __repr__(self): ) def forward(self, x): - """Forward.""" + _is_first_microbatch = self.is_first_microbatch + + # Rewrite quant context + quant_context = _get_fp8_autocast_for_quant_params(self.te_quant_params, self.training) + + with quant_context: + out = self.forward_internal(x, _is_first_microbatch) + + self.is_first_microbatch = False + + return out + def forward_internal(self, x, is_first_microbatch: bool = False): + """Forward.""" if self.config.normalization == "LayerNorm": norm_out = torch.nn.functional.layer_norm( x, [x.size(-1)], self.layer_norm_weight, self.layer_norm_bias, self.eps ) elif self.config.normalization == "RMSNorm": - norm_out = torch.nn.functional.rms_norm(x, [x.size(-1)], self.layer_norm_weight, self.eps) - # weights = [getattr(self, name) for name in self.weight_names] - # weights = torch.cat(weights, dim=0) - weights = self._parameters["weight"] + from primus_turbo.pytorch.ops.normalization import rmsnorm + + norm_out = rmsnorm(x, self.layer_norm_weight, self.eps) + else: + assert False, "Not support normalization type." + + weight = self._parameters["weight"] if self.use_bias: bias_tensor = torch.cat([getattr(self, name) for name in self.bias_names]) + else: + bias_tensor = None + original_shape = x.size() if not norm_out.is_contiguous(): norm_out = norm_out.contiguous() inp = norm_out.view(-1, original_shape[-1]) - if PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled(): - quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() - if quant_config.current_scaling() or quant_config.block_scaling() or quant_config.mxfp8_scaling(): - fp8_gemm = pt.ops.gemm_fp8 - else: - raise ValueError("Not support quant config.") - - out = fp8_gemm( - inp, weights, 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() - if quant_config.mxfp4_scaling(): - fp4_gemm = pt.ops.gemm_fp4 - else: - raise ValueError("Not support quant config.") + if _use_split_wgrad_op(): + from .zbpp_gemm import gemm_with_weight_gradient_store - out = fp4_gemm( - inp, weights, trans_a=False, trans_b=True, out_dtype=None, config=quant_config.data() - ) + out = gemm_with_weight_gradient_store(inp, weight, bias=None) else: - out = self.gemm(inp, weights) + if PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled(): + quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() + assert ( + 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( + weight, + float8_e4m3, + 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_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( + 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(), + ) + else: + out = primus_turbo_torch.ops.gemm(inp, weight, trans_a=False, trans_b=True, out_dtype=None) out = out.view(original_shape[0], original_shape[1], -1) - if self.te_return_bias: - return out, bias_tensor + if self.use_bias: - return out + bias_tensor, None + out = out + bias_tensor + return out, None -class PrimusTurboGroupedMLP(TEGroupedMLP): - """ - Compatibility GroupedMLP for legacy turbo grouped-gemm paths. +def fused_bias_act_with_probs( + intermediate_parallel: torch.Tensor, + bias_parallel: torch.Tensor, + permuted_probs: torch.Tensor, + tokens_per_experts: torch.Tensor, + activation_func: str, +): + assert intermediate_parallel.ndim == 2 + assert permuted_probs.ndim == 1 + assert tokens_per_experts.device == intermediate_parallel.device + + # TODO(ruibin): fuse bias addition with activation function + if bias_parallel is not None: + intermediate_parallel = intermediate_parallel + bias_parallel + + num_tokens = intermediate_parallel.shape[0] + row_mask = primus_turbo_torch.ops.tokens_per_expert_to_mask(tokens_per_experts, num_tokens) + + # TODO(ruibin): support more activation functions + if activation_func == "silu": + fused_act_with_probs = primus_turbo_torch.ops.swiglu_with_probs + elif activation_func == "gelu": + fused_act_with_probs = primus_turbo_torch.ops.geglu_with_probs + else: + raise ValueError(f"Activation function {activation_func} is not supported.") - Megatron removed the old ``GroupedMLP`` implementation, but DeepEP sync-free - stage 2/3 still relies on the turbo grouped-gemm token-count contract. Keep - TEGroupedMLP-style parameter initialization/checkpoint layout while executing - the expert MLP with PrimusTurbo grouped-gemm kernels. + return fused_act_with_probs(intermediate_parallel, permuted_probs, row_mask) + + +class PrimusTurboGroupedLinear(TEGroupedLinear): + """ + Wrapper for the PrimusTurbo `grouped_gemm` ops. """ def __init__( self, - num_local_experts: int, - config: TransformerConfig, - submodules, + num_gemms: int, + input_size: int, + output_size: int, + *, + parallel_mode: Optional[str], + config: ModelParallelConfig, + init_method: Callable, + bias: bool, + skip_bias_add: bool, + is_expert: bool = False, + tp_comm_buffer_name: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, ): args = get_args() - self.offload = args.offload and "grouped_linear" in args.offload_ops - assert not self.offload, "grouped_linear offload is not supported in PrimusTurboGroupedMLP" + self.offload = args.offload and "column_parallel_gemm" in args.offload_ops + assert not self.offload, "gemm offload still have some problems" super().__init__( - num_local_experts=num_local_experts, + num_gemms, + input_size, + output_size, + parallel_mode=parallel_mode, config=config, - submodules=submodules, + init_method=init_method, + bias=bias, + skip_bias_add=skip_bias_add, + is_expert=is_expert, + tp_comm_buffer_name=tp_comm_buffer_name, pg_collection=pg_collection, ) - self.use_turbo_fused_act_with_probs = args.use_turbo_fused_act_with_probs - self.disable_turbo_grouped_mlp_low_precision = args.disable_turbo_grouped_mlp_low_precision - self.patch_zero_bubble = args.patch_zero_bubble - self.patch_primus_pipeline = args.patch_primus_pipeline - if self.config.add_bias_linear: - raise ValueError("PrimusTurboGroupedMLP does not support add_bias_linear=True") + tp_size = get_pg_size(self._tp_group) + assert tp_size == 1, "PrimusTurboGroupedLinear only supports tensor parallel size = 1" - if self.use_turbo_fused_act_with_probs: - assert self.config.gated_linear_unit, "turbo_fused_act_with_probs only support with GLU." + w0 = self.weight0 + buffer = torch.empty( + self.num_gemms, + self.out_features, + self.in_features, + device=w0.device, + dtype=w0.dtype, + ) - if self.config.activation_func == F.silu: - turbo_fused_act_with_probs = pt.ops.swiglu_with_probs - elif self.config.activation_func == F.gelu: - turbo_fused_act_with_probs = pt.ops.geglu_with_probs - else: - raise ValueError("Activation function must be silu or gelu when using GroupedMLP.") - - def _activation_func_with_probs(x, probs, tokens_per_experts): - assert x.ndim == 2 - assert probs.ndim == 1 - num_tokens = x.shape[0] - row_mask = pt.ops.tokens_per_expert_to_mask(tokens_per_experts, num_tokens) - return turbo_fused_act_with_probs(x, probs, row_mask) - - self.activation_func_with_probs = _activation_func_with_probs - - def _stack_grouped_linear_weight(self, module: torch.nn.Module) -> torch.Tensor: - """Stack per-expert ``weight{i}`` tensors into a contiguous ``[E, N, K]`` - buffer for grouped-GEMM consumption. - - Plan-6 P34 routes this through the - :mod:`primus.backends.megatron.core.extensions._triton.stack_grouped_weight` - Triton kernel by default — it collapses the eager - ``torch.stack + transpose(1, 2) + contiguous`` chain (two full - passes over the per-expert weight data, ~289.6 ms / 32 calls in - the plan-5 P32 final EP=8 trace) into a single ``[K, N] -> [N, K]`` - tile-transpose with fused stacking. Set - ``PRIMUS_STACK_GROUPED_WEIGHT_TRITON=0`` to fall back to the - eager chain (kept for A/B testing and as the G37 reference path). - """ - weights = [getattr(module, f"weight{i}") for i in range(self.num_local_experts)] - return stack_grouped_weight(weights) + with torch.no_grad(): + for i in range(self.num_gemms): + weight = getattr(self, f"weight{i}") + buffer[i].copy_(weight) - def forward( + 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)] + + # All experts share the same routing/parallel markers, so weight0's are + # representative for the consolidated parameter. + for attr_name, attr_val in saved_weight_attrs[0].items(): + setattr(self.weights, attr_name, attr_val) + + # Free the per-expert weight{i} Parameters now that their data has been + # consolidated into self.weights. + for i in range(self.num_gemms): + name = f"weight{i}" + if name in self._parameters: + del self._parameters[name] + + gc.collect() + torch.cuda.empty_cache() + + # 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], 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._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 + quant_context = _get_fp8_autocast_for_quant_params(self.te_quant_params, self.training) + + with quant_context: + out = self.forward_internal(x, m_splits, _is_first_microbatch) + + self.is_first_microbatch = False + + return out + + def forward_internal( self, - permuted_local_hidden_states: torch.Tensor, - tokens_per_expert: torch.Tensor, - permuted_probs: torch.Tensor, - routing_map: Optional[torch.Tensor] = None, + x: torch.Tensor, + m_splits: torch.Tensor, + is_first_microbatch: bool = False, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Forward step of the legacy PrimusTurbo grouped-gemm MLP.""" - del routing_map + weights = self.weights + # NOTE: keep x and m_splits on the same device + if m_splits.device != x.device: + m_splits = m_splits.to(x.device) - if self.activation_recompute: - self.activation_checkpoint = tensor_parallel.CheckpointWithoutOutput() - - if self.config.moe_apply_probs_on_input: + if PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled(): + quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() assert ( - self.config.moe_router_topk == 1 - ), "`moe_apply_probs_on_input` only works with `moe_router_topk`=1." - original_dtype = permuted_local_hidden_states.dtype - permuted_local_hidden_states = permuted_probs.unsqueeze(-1) * permuted_local_hidden_states - permuted_local_hidden_states = permuted_local_hidden_states.to(original_dtype) - # Probs already applied, so reset to 1. - permuted_probs = torch.ones_like(permuted_probs) - - w1 = self._stack_grouped_linear_weight(self.linear_fc1) - w2 = self._stack_grouped_linear_weight(self.linear_fc2) - tokens_per_expert = tokens_per_expert.to(w1.device) - - use_grouped_gemm_low_precision = ( - PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled() - and not self.disable_turbo_grouped_mlp_low_precision - ) - probs_for_activation = permuted_probs.unsqueeze(-1) - - if permuted_local_hidden_states.nelement() != 0: - if use_grouped_gemm_low_precision: - quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() - fc1_output = pt.ops.grouped_gemm_fp8( - permuted_local_hidden_states, - w1, - tokens_per_expert, - trans_b=False, - config=quant_config.data(), - ) - else: - fc1_output = pt.ops.grouped_gemm( - permuted_local_hidden_states, w1, tokens_per_expert, trans_b=False + 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( + weights, + float8_e4m3, + quant_config, + disable_parameter_transpose_cache=self.disable_parameter_transpose_cache, ) - if self.activation_recompute: - if self.use_turbo_fused_act_with_probs: - intermediate_parallel = self.activation_checkpoint.checkpoint( - self.activation_func_with_probs, - fc1_output, - permuted_probs, - tokens_per_expert, - ) - else: - intermediate_parallel = self.activation_checkpoint.checkpoint( - self.bias_act_func, fc1_output, None, probs_for_activation - ) - else: - if self.use_turbo_fused_act_with_probs: - intermediate_parallel = self.activation_func_with_probs( - fc1_output, permuted_probs, tokens_per_expert + x, quantized_weights = _bridge_weight_grad( + x, + weights, + PrimusTurboQuantizedTensorPair( + data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer + ), + ) + + # OPT-1 (opt-in, single-GPU): accumulate the expert wgrad straight into + # main_grad in the grouped GEMM backward (beta=1 ACCUMULATE) instead of + # GEMM-wgrad -> _WeightGradBridge.main_grad.add_(). Per-expert shims point + # at the consolidated [E,N,K] main_grad slices; the grouped backward + # accumulates into them and returns grad_b=None, then the bridge backward + # just flags grad_added. ONLY safe with no gradient all-reduce / + # reduce-scatter (TP=1 / DP=1 / EP=1), since grad_b=None skips the reduce. + # Needs a turbo wheel carrying fused_grouped_wgrad (else falls back). + _wgrad_ctx = contextlib.nullcontext() + _fwg_dbg = ( + os.environ.get("PRIMUS_TURBO_FUSE_WGRAD_DEBUG") == "1" + and getattr(type(self), "_fwg_logn", 0) < 10 + ) + _fwg_flag = os.environ.get("PRIMUS_TURBO_FUSE_GROUPED_WGRAD", "0") == "1" + _mg = getattr(weights, "main_grad", None) + if _fwg_flag and _mg is not None and _mg.dim() == 3: + try: + from primus_turbo.pytorch.ops.grouped_gemm_fp8 import ( + _expert_main_grad_view, + fused_grouped_wgrad, ) - else: - intermediate_parallel = self.bias_act_func(fc1_output, None, probs_for_activation) - if use_grouped_gemm_low_precision: - quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() - output = pt.ops.grouped_gemm_fp8( - intermediate_parallel, - w2, - tokens_per_expert, - trans_b=False, + _shims = [_MainGradShim(_mg[i]) for i in range(_mg.shape[0])] + if _fwg_dbg: + import sys + + _v = _expert_main_grad_view(_shims) + print( + f"[OPT-1] gate PASS shape={tuple(_mg.shape)} contig={_mg.is_contiguous()} " + f"stride={_mg.stride()} view={'OK' if _v is not None else 'REJECTED'}", + file=sys.stderr, + flush=True, + ) + type(self)._fwg_logn = getattr(type(self), "_fwg_logn", 0) + 1 + _wgrad_ctx = fused_grouped_wgrad(_shims) + except ImportError as _e: + if _fwg_dbg: + import sys + + print(f"[OPT-1] ImportError: {_e}", file=sys.stderr, flush=True) + type(self)._fwg_logn = getattr(type(self), "_fwg_logn", 0) + 1 + elif _fwg_dbg: + import sys + + print( + f"[OPT-1] gate FAIL flag={_fwg_flag} main_grad={'None' if _mg is None else f'dim={_mg.dim()}'}", + file=sys.stderr, + flush=True, + ) + type(self)._fwg_logn = getattr(type(self), "_fwg_logn", 0) + 1 + + with _wgrad_ctx: + out = primus_turbo_torch.ops.grouped_gemm_fp8( + x, + quantized_weights, + m_splits, + trans_b=True, config=quant_config.data(), ) - else: - output = pt.ops.grouped_gemm(intermediate_parallel, w2, tokens_per_expert, trans_b=False) + elif PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp4_enabled(): + assert False, "FP4 is not supported in PrimusTurboGroupedLinear" else: - # Keep a gradient path for expert weights even when no local token is routed here. - assert ( - not self.patch_zero_bubble and not self.patch_primus_pipeline - ), "Zero bubble or primus pipeline not support torch.matmul backend yet" - w1_flat = w1.view(self.config.hidden_size, -1) - w2_flat = w2.view(-1, self.config.hidden_size) - hidden = torch.matmul(permuted_local_hidden_states, w1_flat) - if self.activation_recompute: - if self.use_turbo_fused_act_with_probs: - hidden = self.activation_checkpoint.checkpoint( - self.activation_func_with_probs, hidden, permuted_probs, tokens_per_expert - ) - else: - hidden = self.activation_checkpoint.checkpoint( - self.bias_act_func, hidden, None, probs_for_activation - ) - else: - if self.use_turbo_fused_act_with_probs: - hidden = self.activation_func_with_probs(hidden, permuted_probs, tokens_per_expert) - else: - hidden = self.bias_act_func(hidden, None, probs_for_activation) - output = torch.matmul(hidden, w2_flat) + out = primus_turbo_torch.ops.grouped_gemm(x, weights, m_splits, trans_b=True) - if self.activation_recompute: - self.activation_checkpoint.discard_output_and_register_recompute(output) + return out, None - return output, None + +class PrimusTurboColumnParallelGroupedLinear(PrimusTurboGroupedLinear): + """ + Wrapper for the PrimusTurboGroupedLinear layer but specialized + to column-parallel style. + """ + + def __init__( + self, + num_gemms: int, + input_size: int, + output_size: int, + *, + config: ModelParallelConfig, + init_method: Callable, + bias: bool, + skip_bias_add: bool, + is_expert: bool, + tp_comm_buffer_name: Optional[str] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + ): + super().__init__( + num_gemms=num_gemms, + input_size=input_size, + output_size=output_size, + parallel_mode="column", + config=config, + init_method=condition_init_method(config, init_method), + bias=bias, + skip_bias_add=skip_bias_add, + is_expert=is_expert, + tp_comm_buffer_name=tp_comm_buffer_name, + pg_collection=pg_collection, + ) + + tp_size = get_pg_size(self._tp_group) + assert tp_size == 1, "PrimusTurboColumnParallelGroupedLinear only supports tensor parallel size = 1" + + +class PrimusTurboRowParallelGroupedLinear(PrimusTurboGroupedLinear): + """ + Wrapper for the PrimusTurboGroupedLinear layer but specialized + to row-parallel style. + """ + + def __init__( + self, + num_gemms: int, + input_size: int, + output_size: int, + *, + config: ModelParallelConfig, + init_method: Callable, + bias: bool, + skip_bias_add: bool, + is_expert: bool, + tp_comm_buffer_name: Optional[str] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + ): + super().__init__( + num_gemms=num_gemms, + input_size=input_size, + output_size=output_size, + parallel_mode="row", + config=config, + init_method=condition_init_method(config, init_method), + bias=bias, + skip_bias_add=skip_bias_add, + is_expert=is_expert, + tp_comm_buffer_name=tp_comm_buffer_name, + pg_collection=pg_collection, + ) + + tp_size = get_pg_size(self._tp_group) + assert tp_size == 1, "PrimusTurboRowParallelGroupedLinear only supports tensor parallel size = 1" class PrimusTurboDeepEPTokenDispatcher(MoETokenDispatcher): @@ -1480,7 +2009,19 @@ def __init__( # fully sync-free moe permute_max_token_num = num_worst_tokens * config.moe_router_topk - self.deepep_dispatcher = pt.modules.DeepEPTokenDispatcher( + pad_multiple = 0 + if args.moe_router_padding_for_quantization: + if PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled(): + pad_multiple = ( + 32 + if PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config().mxfp8_scaling() + else 16 + ) + elif PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp4_enabled(): + pad_multiple = 32 + + use_turbo_grouped_mlp = args.use_turbo_grouped_mlp + self.deepep_dispatcher = primus_turbo_torch.modules.DeepEPTokenDispatcher( num_experts=config.num_moe_experts, router_topk=config.moe_router_topk, ep_group=self.ep_group, @@ -1488,11 +2029,12 @@ def __init__( tp_ep_group=self.tp_ep_group, expert_capacity_factor=config.moe_expert_capacity_factor, permute_fusion=config.moe_permute_fusion, + pad_multiple=pad_multiple, permute_max_token_num=permute_max_token_num, deepep_use_comm_stream=args.turbo_deepep_use_comm_stream, deepep_num_use_cu=args.turbo_deepep_num_cu, deepep_num_worst_tokens=num_worst_tokens, - deepep_use_cuda_num_tokens_per_expert=args.use_turbo_grouped_mlp, + deepep_use_cuda_num_tokens_per_expert=use_turbo_grouped_mlp, deepep_async_finish=True, deepep_allocate_on_comm_stream=True, ) @@ -1635,16 +2177,22 @@ def combine_postprocess(self, hidden_states: torch.Tensor): class PrimusTurboRMSNorm(te.pytorch.RMSNorm): - def __init__(self, *args, **kwargs): - assert "device" in kwargs - assert "dtype" in kwargs or "params_dtype" in kwargs, "device and dtype must be provided" - super().__init__(*args, **kwargs) - self.rms_norm_func = pt.modules.RMSNorm( - normalized_shape=kwargs["hidden_size"], - eps=self.eps, - device=kwargs["device"], - dtype=kwargs["dtype"] if "dtype" in kwargs else kwargs["params_dtype"], + def __init__( + self, + normalized_shape: Union[Iterable[int], int, None] = None, + eps: float = 1e-5, + sequence_parallel: Optional[bool] = None, # legacy + params_dtype: Optional[torch.dtype] = None, # deprecated + zero_centered_gamma: bool = False, + hidden_size: Optional[int] = None, # deprecated + **kwargs, + ): + + super().__init__( + normalized_shape, eps, sequence_parallel, params_dtype, zero_centered_gamma, hidden_size, **kwargs ) def forward(self, x): - return self.rms_norm_func(x) + from primus_turbo.pytorch.ops.normalization import rmsnorm + + return rmsnorm(x, self.weight, self.eps) 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 4bdc0517f..b700bca9c 100644 --- a/primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py +++ b/primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py @@ -39,12 +39,14 @@ from primus.backends.megatron.core.extensions.primus_turbo import ( PrimusTurboAttention, + PrimusTurboColumnParallelGroupedLinear, PrimusTurboColumnParallelLinear, - PrimusTurboGroupedMLP, PrimusTurboLayerNormColumnParallelLinear, PrimusTurboLinear, + PrimusTurboRowParallelGroupedLinear, PrimusTurboRowParallelLinear, ) +from primus.backends.megatron.core.transformer.experts import PrimusGroupedMLP from primus.backends.megatron.training.global_vars import get_primus_args @@ -162,25 +164,27 @@ def grouped_mlp_modules( # and prefer TEGroupedMLP by default. moe_use_legacy_grouped_gemm = getattr(self.cfg, "moe_use_legacy_grouped_gemm", False) - if ( - moe_use_grouped_gemm - and TEColumnParallelGroupedLinear is not None - and not moe_use_legacy_grouped_gemm - ): - _GroupedMLP = ( - PrimusTurboGroupedMLP if getattr(self.cfg, "use_turbo_grouped_mlp", False) else TEGroupedMLP - ) - # TODO: need to update primus_turbo to support TEColumnParallelGroupedLinear? - return _GroupedMLP, TEGroupedMLPSubmodules( - linear_fc1=TEColumnParallelGroupedLinear, linear_fc2=TERowParallelGroupedLinear + use_turbo_grouped_mlp = getattr(self.cfg, "use_turbo_grouped_mlp", False) + if moe_use_grouped_gemm and not moe_use_legacy_grouped_gemm: + _GroupedMLP = PrimusGroupedMLP if use_turbo_grouped_mlp else TEGroupedMLP + GroupedMLPSubmodules = TEGroupedMLPSubmodules( + linear_fc1=( + PrimusTurboColumnParallelGroupedLinear + if use_turbo_grouped_mlp + else TEColumnParallelGroupedLinear + ), + linear_fc2=( + PrimusTurboRowParallelGroupedLinear + if use_turbo_grouped_mlp + else TERowParallelGroupedLinear + ), ) + return _GroupedMLP, GroupedMLPSubmodules elif moe_use_grouped_gemm: warnings.warn( "The legacy GroupedMLP was removed from this Megatron version; " "Primus is using its local compatibility implementation." ) - if getattr(self.cfg, "use_turbo_grouped_mlp", False): - raise NotImplementedError("PrimusTurbo does not support Legacy GroupedMLP") return GroupedMLP, None else: if not is_te_min_version("1.7.0.dev0"): diff --git a/primus/backends/megatron/core/fp8_utils.py b/primus/backends/megatron/core/fp8_utils.py index 27fbd12cc..e3e5b4847 100644 --- a/primus/backends/megatron/core/fp8_utils.py +++ b/primus/backends/megatron/core/fp8_utils.py @@ -34,6 +34,16 @@ # Primus-Turbo not found 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 MX_SCALING_BLOCK_SIZE = 32 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/deepseek_v4/deepseek_v4_block.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py index a007995a2..b06eac1ea 100644 --- a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py +++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py @@ -59,6 +59,7 @@ import ast import logging +from contextlib import nullcontext from dataclasses import dataclass from typing import Callable, List, Optional, Union @@ -999,7 +1000,38 @@ def _recompute_local_layer_indices(self) -> Optional[set]: return set(range(n_local)) raise ValueError(f"Invalid recompute_method for DeepSeek-V4: {method!r}") - def _forward_layer_checkpointed(self, layer, x, position_ids, token_ids): + 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 @@ -1007,11 +1039,16 @@ def _forward_layer_checkpointed(self, layer, x, position_ids, token_ids): 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): - out, _ = layer(hidden, position_ids=position_ids, token_ids=token_ids) + with self._layer_fp8_context(global_idx): + out, _ = layer(hidden, position_ids=position_ids, token_ids=token_ids) return out return tensor_parallel.checkpoint( @@ -1112,14 +1149,16 @@ def forward( # 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) + x = self._forward_layer_checkpointed(layer, x, position_ids, token_ids, global_idx) else: - x, _ = layer( - x, - position_ids=position_ids, - token_ids=token_ids, - ) + 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. 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 index b983820f6..07a9b0013 100644 --- 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 @@ -13,6 +13,7 @@ import importlib import importlib.util import logging +import os from typing import List, Optional, Tuple from megatron.core.transformer.enums import AttnMaskType @@ -130,9 +131,7 @@ def _import_primus_turbo_deepep_dispatcher_cls(): if importlib.util.find_spec("primus_turbo") is None: return None try: - module = importlib.import_module( - "primus.backends.megatron.core.extensions.primus_turbo" - ) + 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 " @@ -220,6 +219,19 @@ 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, @@ -279,7 +291,15 @@ def _build_column_parallel_spec( gather-then-shard variant for full sharded heads is tracked in P14 once the grouped-O TP plan lands. """ - if gather_output: + # 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() @@ -324,7 +344,14 @@ def _build_row_parallel_spec( ``provider.row_parallel_linear_with_scatter_input()``; the standard TE path stays for ``input_is_parallel=True``. """ - if not input_is_parallel: + # 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() 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 index 8662bc8f5..8c0aa7e82 100644 --- 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 @@ -40,9 +40,7 @@ import torch from megatron.core import tensor_parallel -from megatron.core.transformer.multi_token_prediction import ( - MultiTokenPredictionLayer, -) +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 ( @@ -182,9 +180,7 @@ def _proj_and_transformer_layer( # 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 = _lower_streams_out(x, post_process=True, hc_mult=self.hc_mult) # [S, B, D] hidden_states = self._postprocess(hidden_states) return hidden_states 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 index 35a3704a5..b0ae7e583 100644 --- 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 @@ -82,9 +82,9 @@ def _extract_v4_inner_layer_spec(transformer_layer_spec: ModuleSpec) -> ModuleSp 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 - ) + 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: 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 index 6f38ed04d..f3d40315b 100644 --- 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 @@ -111,44 +111,23 @@ class DeepSeekV4TransformerConfig(MLATransformerConfig): # ``PRIMUS_USE_V4_FP8_INDEXER`` (see run_deepseek_v4.sh / flash yaml). use_v4_fp8_indexer: bool = False - # ---- DeepSeek-V4 plan-4 in-tree Triton attention kernels ---- - # Plan-4 P25: when ``use_v4_triton_attention=True`` the dense - # (``compress_ratio == 0``) and HCA (``compress_ratio == 128``) - # paths route through the Primus-owned Triton FlashAttention kernel - # in ``v4_attention_kernels/v4_attention.py`` instead of the - # eager-Python softmax. Precedence is - # ``use_turbo_attention > use_v4_tilelang_attention > use_v4_triton_attention > eager``. - # Plan-4 P26: ``use_v4_triton_csa_attention=True`` enables the CSA - # (``compress_ratio == 4``) Triton kernel; precedence is - # ``use_v4_tilelang_csa_attention > use_v4_triton_csa_attention > eager``. - use_v4_triton_attention: bool = False - use_v4_triton_csa_attention: bool = False - - # ---- DeepSeek-V4 plan-8 tilelang attention kernels (default OFF) ---- - # Plan-8 P57 close-out 2 (2026-05-15): tilelang is an OPTIONAL - # dependency. When the container has tilelang installed at the - # plan-8-pinned version and a caller sets - # ``use_v4_tilelang_attention=True`` (or - # ``use_v4_tilelang_csa_attention=True`` for the CSA family), the - # ``v4_attention`` / ``v4_csa_attention`` dispatcher routes through - # the tilelang FWD/BWD instead of the plan-4 / plan-5 Triton path. - # When tilelang is NOT installed the dispatcher prints a one-time - # rank-0 warning and falls back to Triton -- training continues - # without error. Defaults False so the default container (which - # does not include tilelang) works out of the box. - # - # Replaces the legacy ``PRIMUS_V4_TILELANG_ATTN`` env knob; the env - # var is no longer consulted. See - # ``primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/__init__.py`` - # for the dispatcher; the precedence is: - # 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 - use_v4_tilelang_attention: bool = False - use_v4_tilelang_csa_attention: bool = False + # ---- 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- diff --git a/primus/backends/megatron/core/optimizer/moun.py b/primus/backends/megatron/core/optimizer/moun.py index a40000a32..f14e6a50d 100644 --- a/primus/backends/megatron/core/optimizer/moun.py +++ b/primus/backends/megatron/core/optimizer/moun.py @@ -3,6 +3,7 @@ """Megatron muon optimizer wrapper to handle tensor-parallel.""" import logging +import os from typing import Any, Callable, Dict, List, Literal, Optional import torch @@ -89,8 +90,7 @@ def _resolve_muon_coefficient_type(config, model_chunks) -> str: coeff = str(getattr(config, "muon_coefficient_type", "quintic") or "quintic") if coeff == "quintic": is_v4 = any( - type(getattr(mc, "config", None)).__name__ == "DeepSeekV4TransformerConfig" - for mc in model_chunks + type(getattr(mc, "config", None)).__name__ == "DeepSeekV4TransformerConfig" for mc in model_chunks ) if is_v4: coeff = "deepseekv4" @@ -191,6 +191,34 @@ def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> t # Because -1 is a valid index for ndarray, we decided to not overload it. partition_dim = None + if grad.dim() == 3: + # Consolidated grouped-expert weight [E, N, K]: Muon orthogonalizes 2-D + # matrices, so orthogonalize each expert's [N, K] slice independently + # (matches the per-expert weight{i} path on other branches). partition_dim + # on the 3-D param indexes [E, N, K]; drop the leading expert axis for the + # per-slice call (None under TP=1). + slice_pdim = None if partition_dim is None or partition_dim == 0 else partition_dim - 1 + if slice_pdim is None and os.environ.get("PRIMUS_MUON_BATCHED_NS", "1") != "0": + # Non-partitioned (TP=1) experts: orthogonalize ALL experts in a + # single batched Newton-Schulz instead of an E-long Python loop. + # emerging_optimizers>=0.4.0a0 ``newton_schulz`` accepts 3-D + # ``[E, N, K]`` natively (``batched_newton_schulz_step`` via + # ``torch.baddbmm``), so one call replaces E per-expert calls, + # folding E rounds of Python op-dispatch and the per-expert GEMM + # launches into one batched launch. ``scaled_orthogonalize_fn`` keys + # its scale on ``grad.size(-2)/(-1)`` (= N, K, shared across experts), + # so the single scalar applies to the whole batch -- numerically + # equivalent to the per-slice loop below. + # + # Partitioned experts (TP>1) keep the loop: the batched + + # ``partition_dim`` path is not exercised here. Set + # ``PRIMUS_MUON_BATCHED_NS=0`` to force the per-expert loop. + return self.scaled_orthogonalize_fn(grad, tp_group, None) + return torch.stack( + [self.scaled_orthogonalize_fn(grad[e], tp_group, slice_pdim) for e in range(grad.shape[0])], + dim=0, + ) + 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 diff --git a/primus/backends/megatron/core/transformer/compressor.py b/primus/backends/megatron/core/transformer/compressor.py index 1c079a71e..c99679edd 100644 --- a/primus/backends/megatron/core/transformer/compressor.py +++ b/primus/backends/megatron/core/transformer/compressor.py @@ -28,6 +28,7 @@ from __future__ import annotations +import os from typing import Optional import torch @@ -77,8 +78,16 @@ def __init__( self.coff = 2 if self.overlap else 1 proj_out = self.coff * head_dim - self.wkv = nn.Linear(hidden_size, proj_out, bias=False) - self.wgate = nn.Linear(hidden_size, proj_out, bias=False) + self._proj_out = proj_out + # 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 @@ -90,6 +99,21 @@ def __init__( 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 # ------------------------------------------------------------------ @@ -127,8 +151,11 @@ def _overlap_transform(self, t: torch.Tensor) -> torch.Tensor: def forward(self, hidden: torch.Tensor) -> torch.Tensor: """Pool ``hidden[B, S, D]`` to ``[B, S/ratio, head_dim]``.""" - kv_proj = self.wkv(hidden) # [B, S, coff*head_dim] - score_proj = self.wgate(hidden) # [B, S, coff*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] @@ -138,13 +165,31 @@ def forward(self, hidden: torch.Tensor) -> torch.Tensor: score = self._overlap_transform(score) # [B, N, 2*ratio, head_dim] # else: kv / score already at [B, N, ratio, head_dim] - # APE adds a per-window-slot bias to the score (broadcast over B, N). - score = score + self.ape # [B, N, win, head_dim] - - # Softmax over the per-window axis (dim=2). Each compressed token is - # a weighted average of its window members. - weights = F.softmax(score.float(), dim=2).to(kv.dtype) - pooled = (kv * weights).sum(dim=2) # [B, N, 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 diff --git a/primus/backends/megatron/core/transformer/deepseek_v4_attention.py b/primus/backends/megatron/core/transformer/deepseek_v4_attention.py index 9842923a0..543f7fbd9 100644 --- a/primus/backends/megatron/core/transformer/deepseek_v4_attention.py +++ b/primus/backends/megatron/core/transformer/deepseek_v4_attention.py @@ -81,11 +81,28 @@ 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_ +# (dense/HCA) and v4_csa_attention_ (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, - v4_attention, - v4_csa_attention_from_pool, + 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, ) _SUPPORTED_COMPRESS_RATIOS = (0, 4, 128) @@ -93,8 +110,30 @@ 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 +# P32 diagnostic: collect in-context cuda.Event timings of v4_attention_v1 # --------------------------------------------------------------------------- @@ -123,7 +162,7 @@ def dump(cls) -> None: local_rank = 0 if local_rank != 0: return - print("[PRIMUS_V4_DIAG_TIME] v4_attention inline cuda.Event timings:", flush=True) + 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 @@ -241,6 +280,55 @@ def _projection_forward(proj: nn.Module, x: torch.Tensor) -> torch.Tensor: 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``. @@ -287,11 +375,15 @@ def _per_head_rms_norm(x: torch.Tensor, *, eps: float) -> torch.Tensor: 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. """ - in_dtype = x.dtype - x32 = x.float() - rsqrt = torch.rsqrt(x32.square().mean(dim=-1, keepdim=True) + eps) - return (x32 * rsqrt).to(in_dtype) + 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: @@ -349,11 +441,11 @@ class DeepseekV4Attention(MLASelfAttention): compress_ratio == 0 (dense / SWA, single key axis): use_turbo_attention > use_v4_triton_attention > eager - (-> self.core_attention) (-> v4_attention) (-> _attention_forward) + (-> self.core_attention) (-> v4_attention_v1) (-> _attention_forward) compress_ratio == 128 (HCA: local SWA + full compressed pool): use_v4_triton_attention > eager - (-> v4_attention, (-> _attention_forward + (-> v4_attention_v1, (-> _attention_forward HCA path with joint with cat([local, pool]) [local | pool] mask) additive mask) @@ -363,7 +455,7 @@ class DeepseekV4Attention(MLASelfAttention): compress_ratio == 4 (CSA: local SWA + per-query top-K gather): use_v4_triton_csa_attention > eager - (-> v4_csa_attention) (-> _csa_forward eager) + (-> v4_csa_attention_v0) (-> _csa_forward eager) Neither ``use_turbo_attention`` nor ``use_v4_triton_attention`` applies to CSA — the per-query @@ -380,7 +472,7 @@ class DeepseekV4Attention(MLASelfAttention): 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 + 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. """ @@ -595,62 +687,88 @@ def __init__( # 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``. - self._use_v4_triton_attention: bool = bool(getattr(config, "use_v4_triton_attention", False)) - if self._use_v4_triton_attention and self.compress_ratio not in (0, 128): - # Plan-4 P26 ships the CSA Triton kernel under - # ``use_v4_triton_csa_attention`` — the dense / HCA flag does - # NOT enable it. Surface the misconfiguration loud at build - # time so a stray run script doesn't silently fall back to - # eager for the CSA layers (and skew layer-vs-layer perf). - self._use_v4_triton_attention = False - - # Plan-4 P26: in-tree Primus Triton kernel for cr == 4 (CSA). - # Symmetric to ``_use_v4_triton_attention`` above: read the flag - # once at __init__, and auto-disable for non-CSA layers so a - # stray ``use_v4_triton_csa_attention=True`` does not silently - # accelerate the CSA layers only and skew apples-to-apples perf - # comparisons. Precedence in ``forward`` (cr == 4 branch) is - # ``use_v4_triton_csa_attention > eager``. - self._use_v4_triton_csa_attention: bool = bool(getattr(config, "use_v4_triton_csa_attention", False)) - if self._use_v4_triton_csa_attention and self.compress_ratio != 4: - self._use_v4_triton_csa_attention = False - - # Plan-8 P57 close-out 2: optional tilelang dispatch flags - # (replace the legacy PRIMUS_V4_TILELANG_ATTN env knob). Each - # flag is layer-kind specific: - # - # - ``use_v4_tilelang_attention`` -> cr ∈ {0, 128} - # auto-disabled at non-dense/HCA layers symmetric to the - # ``use_v4_triton_attention`` rule. - # - ``use_v4_tilelang_csa_attention`` -> cr == 4 - # auto-disabled at non-CSA layers symmetric to the - # ``use_v4_triton_csa_attention`` rule. - # - # When either flag is set but tilelang is not installed (or the - # plan-8 P50..P55 kernels are not registered) the dispatcher - # falls back to the Triton path with a one-time rank-0 warning; - # no runtime error is raised so the default container (which - # does NOT ship tilelang) just runs Triton transparently. - # - # We use a string-aware boolean coercion because the yaml - # default ``${PRIMUS_USE_V4_TILELANG_ATTENTION:false}`` resolves - # to the STRING ``"false"`` when the env var is unset, and - # ``bool("false") is True``. Existing flags - # (``use_v4_triton_attention`` etc.) do not trip this because - # the V4 run scripts always pass them via ``-- "False"`` - # CLI which the override parser converts to a Python ``False``. - self._use_v4_tilelang_attention: bool = _coerce_optional_bool_flag( - getattr(config, "use_v4_tilelang_attention", False), - field_name="use_v4_tilelang_attention", + # ---- 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", ) - if self._use_v4_tilelang_attention and self.compress_ratio not in (0, 128): - self._use_v4_tilelang_attention = False - self._use_v4_tilelang_csa_attention: bool = _coerce_optional_bool_flag( - getattr(config, "use_v4_tilelang_csa_attention", False), - field_name="use_v4_tilelang_csa_attention", + 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._use_v4_tilelang_csa_attention and self.compress_ratio != 4: - self._use_v4_tilelang_csa_attention = False + 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 @@ -743,29 +861,12 @@ def _log_kernel_choice(self) -> None: if self.compress_ratio == 0: if self._use_core_attention: kernel = "core_attention (Turbo / TE flash)" - elif self._use_v4_triton_attention: - if self._use_v4_tilelang_attention: - kernel = "v4_attention (tilelang->Triton fallback, dense path)" - else: - kernel = "v4_attention (Triton, dense path)" else: - kernel = "v4_attention (eager Python, dense path)" + kernel = f"dense attention backend = {self._attn_backend}" elif self.compress_ratio == 128: - if self._use_v4_triton_attention: - if self._use_v4_tilelang_attention: - kernel = "v4_attention (tilelang->Triton fallback, HCA path)" - else: - kernel = "v4_attention (Triton, HCA path)" - else: - kernel = "v4_attention (eager Python, HCA path)" + kernel = f"HCA attention backend = {self._attn_backend}" elif self.compress_ratio == 4: - if self._use_v4_triton_csa_attention: - if self._use_v4_tilelang_csa_attention: - kernel = "v4_csa_attention_from_pool (tilelang->Triton fallback)" - else: - kernel = "v4_csa_attention_from_pool (Triton)" - else: - kernel = "v4_csa_attention (eager Python)" + kernel = f"CSA attention backend = {self._csa_backend}" else: # Defensive: __init__ already raises ValueError for unsupported # compress_ratio, so this branch should be unreachable. @@ -912,7 +1013,7 @@ def _attention_forward( for the dense / HCA paths (single key axis). Plan-4 P24: math lives in - :func:`primus...v4_attention_kernels.reference.eager_v4_attention` + :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 @@ -944,7 +1045,7 @@ def _attention_forward_via_v4_triton( hca_local_seqlen: int = 0, ) -> torch.Tensor: """Run the dense / HCA softmax-and-attend through the plan-4 - in-tree :func:`v4_attention` Triton kernel. + 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 @@ -983,7 +1084,7 @@ def _attention_forward_via_v4_triton( ev_s = torch.cuda.Event(enable_timing=True) ev_e = torch.cuda.Event(enable_timing=True) ev_s.record() - out = v4_attention( + out = v4_attention_v1( q, k, v, @@ -994,13 +1095,12 @@ def _attention_forward_via_v4_triton( training=self.training, scale=self._attention_scale(), hca_local_seqlen=int(hca_local_seqlen), - use_tilelang=self._use_v4_tilelang_attention, ) 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( + return v4_attention_v1( q, k, v, @@ -1011,7 +1111,6 @@ def _attention_forward_via_v4_triton( training=self.training, scale=self._attention_scale(), hca_local_seqlen=int(hca_local_seqlen), - use_tilelang=self._use_v4_tilelang_attention, ) def _attention_forward_via_core( @@ -1090,7 +1189,10 @@ def _grouped_o_projection(self, attn: torch.Tensor) -> torch.Tensor: o = o.view(B, S, G * self.o_lora_rank) else: wo_a_w = weight.view(G, self.o_lora_rank, (H * Dh) // G) - o = torch.einsum("bsgd,grd->bsgr", attn_g, wo_a_w) + 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) @@ -1112,9 +1214,10 @@ def _build_compressed_pool(self, hidden: torch.Tensor) -> torch.Tensor: pooled = self.compressor(hidden) # [B, P, head_dim] B, P = pooled.shape[0], pooled.shape[1] - # Compress-base partial RoPE on compressed indices [0..P). - comp_pos = torch.arange(P, device=device) - cos, sin = self.rope.compress_rope(comp_pos) + # Compress-base partial RoPE on compressed indices [0..P). Positions are + # the deterministic arange(P), so use the cached table instead of + # rebuilding arange -> outer -> cos/sin every forward. + cos, sin = self.rope.compress_rope.forward_arange(P, device) cos = cos[..., : self.rotary_dim // 2] sin = sin[..., : self.rotary_dim // 2] cos = cos.unsqueeze(0).expand(B, -1, -1) @@ -1148,11 +1251,34 @@ def _hca_extra_kv( # Move heads dim to dim=1: [B, H, P, head_dim]. pool_bh = pool_h.transpose(1, 2) - 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] - extra_mask = torch.where(s_end <= t, 0.0, float("-inf")).to(dtype) + 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, @@ -1171,7 +1297,7 @@ def _csa_forward( 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.reference.eager_v4_csa_attention` + :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 @@ -1195,11 +1321,90 @@ def _csa_forward( # 2) Indexer top-K per query. topk_idxs, _ = self.indexer(hidden) # [B, S, K] - - # Plan-5 P31 dispatch: Triton path consumes pool + topk directly - # and avoids materialising [B, S, K, Dh] gathered tensors. - if self._use_v4_triton_csa_attention: - return v4_csa_attention_from_pool( + # 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, @@ -1210,12 +1415,9 @@ def _csa_forward( attn_dropout=self.attn_dropout, training=self.training, scale=self._attention_scale(), - use_tilelang=self._use_v4_tilelang_csa_attention, ) - # 3) Eager fallback gathers per-query pool slices: [B, S, K, Dh]. - # ``gathered`` is broadcast across heads in the reference op (no - # H dim), matching V4's single-latent pool shared by all heads. + # 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) @@ -1225,6 +1427,20 @@ def _csa_forward( 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, @@ -1242,6 +1458,101 @@ def _csa_forward( # 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, @@ -1287,51 +1598,33 @@ def forward( v_local_bh = v_h.transpose(1, 2) if self.compress_ratio == 0: - # Plan-4 P25: ``use_v4_triton_attention`` routes the dense - # softmax-and-attend through the in-tree Triton kernel. - # Precedence ``use_turbo_attention > use_v4_triton_attention - # > eager`` is enforced by the earlier ``_use_core_attention`` - # branch returning before this block. - if self._use_v4_triton_attention: - out_bh = self._attention_forward_via_v4_triton( - q_bh, - k_local_bh, - v_local_bh, - None, - swa_window=int(self.attn_sliding_window), - ) - else: - # Eager-Python dense path (used when ``core_attention`` is not - # built or the V4 sink + Turbo sink-attention contract isn't - # met; e.g. CPU unit tests or TE-without-sink configs). - local_mask = self._local_mask(S, device=device, dtype=dtype) - out_bh = self._attention_forward(q_bh, k_local_bh, v_local_bh, local_mask) + 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 cannot use ``core_attention``: the local SWA branch and - # the compressed-pool branch share **one** softmax with **one** - # sink column. Stock flash-attn returns no LSE, so we can't - # decompose into two flash calls and recombine. Plan-4 P25's - # in-tree Triton kernel is HCA-aware (it consumes the joint - # ``cat([local_mask, pool_mask])`` additive bias and runs a - # single online-softmax pass), so HCA opts into - # ``use_v4_triton_attention`` too — exactly the same kernel - # call as the dense path. + # 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) - if self._use_v4_triton_attention: - out_bh = self._attention_forward_via_v4_triton( - q_bh, - k_full, - v_full, - extra_mask, - swa_window=int(self.attn_sliding_window), - hca_local_seqlen=S, - ) - else: - local_mask = self._local_mask(S, device=device, dtype=dtype) - full_mask = torch.cat([local_mask, extra_mask], dim=-1) # [S, S+P] - out_bh = self._attention_forward(q_bh, k_full, v_full, full_mask) + 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 diff --git a/primus/backends/megatron/core/transformer/dual_rope.py b/primus/backends/megatron/core/transformer/dual_rope.py index f93065ab0..4744b889d 100644 --- a/primus/backends/megatron/core/transformer/dual_rope.py +++ b/primus/backends/megatron/core/transformer/dual_rope.py @@ -42,8 +42,9 @@ import torch import torch.nn as nn -from primus.backends.megatron.core.transformer.v4_attention_kernels._triton.rope_interleaved_partial import ( +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rope_interleaved_partial import ( RoPEInterleavedPartialFn, + apply_rope_from_positions, ) # --------------------------------------------------------------------------- @@ -148,6 +149,10 @@ def __init__( # YaRN's m_scale; exposed for the caller to multiply attention scale. self.attn_scale: float = _yarn_attn_scale(self.yarn_factor) + # Memo of (cos, sin) tables keyed by (n, device) for arange(n) 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. @@ -163,6 +168,26 @@ def forward(self, position_ids: torch.Tensor) -> tuple: sin = freqs.sin() return cos, sin + def forward_arange(self, n: int, device) -> tuple: + """``(cos, sin)`` for positions ``torch.arange(n)`` -- cached. + + Equivalent to ``self.forward(torch.arange(n, device=device))``, but + memoised by ``(n, device)``. The compressed-branch RoPE is always + evaluated at the deterministic positions ``arange(P)`` + (``P = S // compress_ratio``, 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). Set + ``PRIMUS_COMPRESS_ROPE_CACHE=0`` to disable the cache. + """ + if os.environ.get("PRIMUS_COMPRESS_ROPE_CACHE", "1") == "0": + return self.forward(torch.arange(n, device=device)) + key = (int(n), str(device)) + hit = self._arange_cache.get(key) + if hit is None: + hit = self.forward(torch.arange(n, device=device)) + self._arange_cache[key] = hit + return hit + # --------------------------------------------------------------------------- # Partial interleaved RoPE application @@ -306,14 +331,23 @@ def apply_rope( position_ids: torch.Tensor, compress_ratio: int, ) -> torch.Tensor: - """Convenience: pick the right rope, build cos/sin, apply partial RoPE. - - ``position_ids`` shape ``[B, S]`` or ``[S]``. Cos/sin will broadcast - over any extra heads dim of ``x``. + """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) - cos, sin = rope(position_ids) - return apply_interleaved_partial_rope(x, cos, sin, rotary_dim=self.rotary_dim) + return apply_rope_from_positions(x, position_ids, rope.inv_freq, rotary_dim=self.rotary_dim) # Convenience accessors for callers who need the YaRN m_scale (e.g. to # adjust attention softmax scale on compressed layers). diff --git a/primus/backends/megatron/core/transformer/experts.py b/primus/backends/megatron/core/transformer/experts.py new file mode 100644 index 000000000..694db2137 --- /dev/null +++ b/primus/backends/megatron/core/transformer/experts.py @@ -0,0 +1,238 @@ +from __future__ import annotations + +from typing import Optional, Tuple, Union + +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, +) +from megatron.core.transformer.moe.experts import TEGroupedMLP, TEGroupedMLPSubmodules +from megatron.core.transformer.moe.moe_utils import ProcessGroupCollection +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.typed_torch import apply_module +from megatron.training.global_vars import get_args + + +class PrimusGroupedMLP(TEGroupedMLP): + """An efficient implementation of the Experts layer using TE's GroupedLinear. + + Executes multiple experts in parallel to maximize computational efficiency. + """ + + def __init__( + self, + num_local_experts: int, + config: TransformerConfig, + submodules: TEGroupedMLPSubmodules, + pg_collection: Optional[ProcessGroupCollection] = None, + ): + args = get_args() + + super().__init__( + num_local_experts, + config, + submodules, + pg_collection, + ) + + # 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 + + 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, + intermediate_parallel: torch.Tensor, + bias_parallel: torch.Tensor, + permuted_probs: torch.Tensor, + tokens_per_experts: Union[torch.Tensor, None] = None, + ): + if self.use_turbo_fused_act_with_probs: + from primus.backends.megatron.core.extensions.primus_turbo import ( + fused_bias_act_with_probs, + ) + + assert ( + tokens_per_experts is not None + ), "tokens_per_experts is required when `use_turbo_fused_act_with_probs` is True." + + if self.activation_func == F.silu and self.config.gated_linear_unit: + activation = "silu" + elif self.activation_func == F.gelu and self.config.gated_linear_unit: + activation = "gelu" + else: + raise ValueError( + "Only support fusion of swiglu and gelu in PrimusGroupedMLP when `use_turbo_fused_act_with_probs` is True." + ) + + # `forward()` unsqueeze(-1)'s `permuted_probs` to [tokens, 1] so the non-fused + # asserts ndim == 1. Squeeze back to 1D for the fused kernel only. + probs_1d = permuted_probs.squeeze(-1) if permuted_probs.dim() == 2 else permuted_probs + # dtype is handled inside the fused kernel + return fused_bias_act_with_probs( + intermediate_parallel, bias_parallel, probs_1d, tokens_per_experts, activation + ) + else: + # 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, + tokens_per_expert: torch.Tensor, + permuted_probs: torch.Tensor, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Forward of PrimusGroupedMLP + + Args: + permuted_local_hidden_states (torch.Tensor): The permuted input hidden states of the + local experts. + tokens_per_expert (torch.Tensor): The number of tokens per expert. + permuted_probs (torch.Tensor): The permuted probs of each token produced by the router. + + 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): + # 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 + permuted_local_hidden_states, tokens_per_expert_cpu = self.quantization_padding( + permuted_local_hidden_states, tokens_per_expert_cpu + ) + permuted_probs, _ = self.quantization_padding( + permuted_probs.unsqueeze(-1), actual_tokens_per_expert_cpu + ) + tokens_per_expert = torch.tensor( + tokens_per_expert_cpu, device=permuted_local_hidden_states.device + ) + else: + permuted_probs = permuted_probs.unsqueeze(-1) + + if self.config.moe_apply_probs_on_input: + assert ( + self.config.moe_router_topk == 1 + ), "`moe_apply_probs_on_input` only works with `moe_router_topk`=1." + original_dtype = permuted_local_hidden_states.dtype + permuted_local_hidden_states = permuted_probs * permuted_local_hidden_states + permuted_local_hidden_states = permuted_local_hidden_states.to(original_dtype) + # Probs already applied, so reset to 1. + permuted_probs = torch.ones_like(permuted_probs) + + with off_interface( + self.offload_expert_fc1, permuted_local_hidden_states, "expert_fc1" + ) as permuted_local_hidden_states: + fc1_output, bias_parallel = apply_module(self.linear_fc1)( + permuted_local_hidden_states, tokens_per_expert + ) + if self.offload_expert_fc1: + fc1_output = off_interface.group_commit( + fc1_output, + name="expert_fc1", + forced_released_tensors=[permuted_local_hidden_states], + ) + + if self.activation_recompute: + self.activation_checkpoint = tensor_parallel.CheckpointWithoutOutput() + with off_interface(self.offload_moe_act, fc1_output, "moe_act") as fc1_output: + # NOTE: use the bias_act_func_with_mask instead of the bias_act_func to reduce the extra compute when stage of `sync_free_moe` is 3. + bias_act_output = self.activation_checkpoint.checkpoint( + self.bias_act_func_with_mask, fc1_output, bias_parallel, permuted_probs, tokens_per_expert + ) + 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) + 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) + + # Delay the offload of the moe act until after the linear_fc2 has been computed + # 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]) + 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): + output = self.quantization_unpadding(output, actual_tokens_per_expert_cpu) + + output_bias = None + + return output, output_bias diff --git a/primus/backends/megatron/core/transformer/hyper_connection.py b/primus/backends/megatron/core/transformer/hyper_connection.py index 220576ac7..d3020e3f0 100644 --- a/primus/backends/megatron/core/transformer/hyper_connection.py +++ b/primus/backends/megatron/core/transformer/hyper_connection.py @@ -55,16 +55,31 @@ import torch.nn as nn import torch.nn.functional as F -from primus.backends.megatron.core.transformer.v4_attention_kernels._triton.hc_glue import ( +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.hc_glue import ( +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.hc_glue import ( +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.sinkhorn import ( +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, @@ -194,7 +209,7 @@ def sinkhorn_normalize( 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.sinkhorn.SinkhornNormalizeFn`. + :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 @@ -297,11 +312,16 @@ def _packed_logits(self, x: torch.Tensor) -> torch.Tensor: ``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) - flat32 = flat.float() - rsqrt = torch.rsqrt(flat32.pow(2).mean(dim=-1, keepdim=True) + self.eps) - logits = F.linear(flat32 * rsqrt, self.fn.weight.to(dtype=flat32.dtype)) + 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 ------------------------------------------------------ @@ -355,8 +375,16 @@ def collapse(x: torch.Tensor, pre: torch.Tensor) -> torch.Tensor: ``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 (pre.unsqueeze(-1) * x).sum(dim=-2) + return hc_collapse_triton(x, pre) @staticmethod def expand( @@ -377,6 +405,9 @@ def expand( 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] @@ -417,10 +448,15 @@ def reset_parameters(self) -> None: nn.init.ones_(self.scale) def forward(self, x: torch.Tensor) -> torch.Tensor: - """``x``: ``[..., K, D]`` → ``[..., D]``.""" - flat = x.flatten(-2).float() - rsqrt = torch.rsqrt(flat.pow(2).mean(dim=-1, keepdim=True) + self.eps) - mixes = F.linear(flat * rsqrt, self.fn.weight.to(dtype=flat.dtype)) # [..., K] + """``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) diff --git a/primus/backends/megatron/core/transformer/indexer.py b/primus/backends/megatron/core/transformer/indexer.py index e5a8284e0..bc06ca642 100644 --- a/primus/backends/megatron/core/transformer/indexer.py +++ b/primus/backends/megatron/core/transformer/indexer.py @@ -45,6 +45,7 @@ from __future__ import annotations import logging +import os from typing import Tuple import torch @@ -62,6 +63,7 @@ def _is_rank0() -> bool: pass return True + from primus.backends.megatron.core.transformer.compressor import Compressor # E4M3 finite max magnitude (float8_e4m3fn): largest representable value. @@ -88,25 +90,98 @@ def fake_quantize_fp8_e4m3(x: torch.Tensor) -> torch.Tensor: 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.indexer_score import ( + + +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.indexer_score import ( +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.indexer_score import ( +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.indexer_score_post import ( +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.indexer_score_post import ( +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.indexer_score_post import ( +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). + + Reuses the attention-proj flag PRIMUS_V4_FP8_ATTN_PROJ; 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_ATTN_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. @@ -154,12 +229,20 @@ def __init__( "BF16 index-score + top-k preserved)." ) - # W^{DQ}: low-rank query down-projection. - self.w_dq = nn.Linear(hidden_size, self.dq_rank, bias=False) + # 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) - # W^w_h: per-head scalar weight. - self.w_w = nn.Linear(hidden_size, index_n_heads, bias=False) # Mini-Compressor producing K^{IComp}. self.indexer_compressor = Compressor( @@ -168,6 +251,21 @@ def __init__( ratio=compress_ratio, ) + 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( @@ -184,10 +282,24 @@ def _causal_mask( 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 # ------------------------------------------------------------------ @@ -219,9 +331,18 @@ def forward(self, hidden: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: k_icomp = k_icomp.unsqueeze(2) # [B, P, 1, Hd] # 2) Per-head query and per-head weight. - q_q = self.w_dq(hidden) # [B, S, dq_rank] - q_i = self.w_iuq(q_q).view(B, S, H, Hd) # [B, S, H, Hd] - w_i = self.w_w(hidden) # [B, S, H] + # Indexer projections: FP8 (paper / NVIDIA backend.linear) when enabled, + # else the bf16 nn.Linear. No TP gather/scatter (duplicated linears). + # FP8 (paper / NVIDIA backend.linear) when enabled, else the bf16 nn.Linear. + 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] # 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 @@ -235,18 +356,25 @@ def forward(self, hidden: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: # else → fully eager. k_icomp_2d = k_icomp.squeeze(2) - # FP8 (E4M3) QK path: fake-quantize the query / compressed-key - # activations BEFORE the scoring einsum so every dispatch path - # (full-fuse Triton, post-einsum-tail Triton, eager) consumes the - # FP8-rounded values. The ReLU + per-head weight + sum + causal mask + - # top-k stay in the activation dtype (BF16 index-score path). ``w_i`` - # (per-head score weights) is left in BF16 — it is not part of the QK - # GEMM the report quantizes. - if self.use_fp8_qk: + # 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) - if _indexer_triton_full_enabled() and _indexer_triton_full_supported(q_i, k_icomp_2d, w_i): + # 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, diff --git a/primus/backends/megatron/core/transformer/local_rmsnorm.py b/primus/backends/megatron/core/transformer/local_rmsnorm.py index 402e544a8..6efa92dad 100644 --- a/primus/backends/megatron/core/transformer/local_rmsnorm.py +++ b/primus/backends/megatron/core/transformer/local_rmsnorm.py @@ -37,6 +37,10 @@ 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. @@ -78,10 +82,22 @@ def __init__( self.eps = float(eps) def forward(self, x: torch.Tensor) -> torch.Tensor: - in_dtype = x.dtype - x32 = x.float() - rsqrt = torch.rsqrt(x32.pow(2).mean(dim=-1, keepdim=True) + self.eps) - return (x32 * rsqrt).to(in_dtype) * self.weight + # 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/backends/megatron/core/transformer/moe/_triton/v4_router_post.py b/primus/backends/megatron/core/transformer/moe/_triton/v4_router_post.py index 66220c65f..82a40c2fa 100644 --- a/primus/backends/megatron/core/transformer/moe/_triton/v4_router_post.py +++ b/primus/backends/megatron/core/transformer/moe/_triton/v4_router_post.py @@ -78,47 +78,55 @@ @triton.jit def _v4_router_post_fwd_kernel( LOGITS_PTR, # [N, E] fp32 - INDICES_PTR, # [N, K] int64 (gather positions) + 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, K] fp32 (saved-for-backward; gathered weights, post-denom, pre-scale) + 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: tl.constexpr, - K: tl.constexpr, + 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. - Layout: per-row state in fp32 registers (`[BLOCK_N, E]` scores - + `[BLOCK_N, K]` weights). At V4-Flash N=4096, E=256, K=6 the - per-row footprint is ~1 KiB which lets us pick BLOCK_N=16 without - hitting register pressure. + 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, E) - k_idx = tl.arange(0, K) + 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, E] in fp32. + # 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], + mask=n_mask[:, None] & e_mask[None, :], other=0.0, ).to(tl.float32) - # Apply score_fn (compile-time specialised). + # 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 - m = tl.max(logits, axis=1, keep_dims=True) - e = tl.exp(logits - m) + 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 @@ -128,33 +136,39 @@ def _v4_router_post_fwd_kernel( # Stable: for large x, log(1+exp(x)) ≈ x. scores = tl.sqrt(softplus) - # Save full scores row for backward. + # 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], + mask=n_mask[:, None] & e_mask[None, :], ) - # Load indices [BLOCK_N, K] (int64). + # Load indices [BLOCK_N, BLOCK_K] (int64; padded cols read as 0, masked below). indices = tl.load( - INDICES_PTR + n_offs[:, None] * K + k_idx[None, :], - mask=n_mask[:, None], + 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 by computing - # the gather entirely in registers (avoid store-then-reload-via-HBM - # hazard). Build weights[BLOCK_N, K] as the per-(n, k) score at - # index indices[n, k] by static loop over K: - weights = tl.zeros((BLOCK_N, K), dtype=tl.float32) - for k_off in tl.static_range(K): - idx_col = tl.load(INDICES_PTR + n_offs * K + k_off, mask=n_mask, 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) - k_is_off = (k_idx == k_off).to(tl.float32) 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) @@ -170,41 +184,44 @@ def _v4_router_post_fwd_kernel( # Save weights for backward. tl.store( - WEIGHTS_OUT_PTR + n_offs[:, None] * K + k_idx[None, :], + WEIGHTS_OUT_PTR + n_offs[:, None] * BLOCK_K + k_idx[None, :], weights, - mask=n_mask[:, None], + 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). + # 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], + mask=n_mask[:, None] & k_mask[None, :], ) tl.store( RMAP_PTR + n_offs[:, None] * E + indices, - tl.full((BLOCK_N, K), 1, dtype=tl.int1), - mask=n_mask[:, None], + 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, K] int64 + INDICES_PTR, # [N, BLOCK_K] int64 (padded) SCORES_PTR, # [N, E] fp32 saved - WEIGHTS_PTR, # [N, K] fp32 saved (post-scaled) + WEIGHTS_PTR, # [N, BLOCK_K] fp32 saved (post-scaled, padded) DENOM_PTR, # [N] fp32 saved DLOGITS_PTR, # [N, E] fp32 OUT N, - E: tl.constexpr, - K: tl.constexpr, + 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. @@ -239,26 +256,29 @@ def _v4_router_post_bwd_kernel( n_offs = pid * BLOCK_N + tl.arange(0, BLOCK_N) n_mask = n_offs < N - e_idx = tl.arange(0, E) - k_idx = tl.arange(0, K) + 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], + mask=n_mask[:, None] & e_mask[None, :], other=0.0, ) indices = tl.load( - INDICES_PTR + n_offs[:, None] * K + k_idx[None, :], - mask=n_mask[:, None], + 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, K], fp32. + # 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], + mask=n_mask[:, None] & k_mask[None, :], other=0.0, ).to(tl.float32) # Chain through the SCALE multiply: forward was @@ -276,8 +296,8 @@ def _v4_router_post_bwd_kernel( # 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] * K + k_idx[None, :], - mask=n_mask[:, None], + 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) @@ -287,15 +307,18 @@ def _v4_router_post_bwd_kernel( # 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 K. For each k, we add d_gathered[:, k] at the position + # 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. - d_scores_full = tl.zeros((BLOCK_N, E), dtype=tl.float32) - for k_off in tl.static_range(K): - idx_col = tl.load(INDICES_PTR + n_offs * K + k_off, mask=n_mask, other=0) - grad_col = tl.load(INDICES_PTR + n_offs * K + k_off, mask=n_mask, other=0) # placeholder - _ = grad_col # unused; the real grad comes from d_gathered slice + # 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] @@ -323,7 +346,7 @@ def _v4_router_post_bwd_kernel( tl.store( DLOGITS_PTR + n_offs[:, None] * E + e_idx[None, :], d_logits, - mask=n_mask[:, None], + mask=n_mask[:, None] & e_mask[None, :], ) @@ -366,26 +389,34 @@ def forward( # type: ignore[override] ) N, E = logits.shape _, K = indices.shape - if E & (E - 1) != 0: - raise ValueError(f"E must be a power of 2; got E={E}") - if K & (K - 1) != 0: - raise ValueError(f"K must be a power of 2; got K={K}") + # 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.empty((N, K), 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 ~ E + K + few scalars in fp32. - if E <= 64: + # BLOCK_N heuristic: per-row state ~ BLOCK_E + BLOCK_K + few scalars. + if block_e <= 64: block_n = 64 - elif E <= 256: + elif block_e <= 256: block_n = 16 else: block_n = 4 @@ -393,19 +424,21 @@ def forward( # type: ignore[override] _v4_router_post_fwd_kernel[grid]( logits_c, - indices_c, + indices_pad, probs, routing_map, scores_saved, weights_saved, denom_saved, N, - E=E, - K=K, + 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, @@ -414,27 +447,31 @@ def forward( # type: ignore[override] }[out_dtype], ) - ctx.save_for_backward(indices_c, scores_saved, weights_saved, denom_saved) + 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_c, scores_saved, weights_saved, denom_saved = ctx.saved_tensors + 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_c.shape[0] - device = indices_c.device + N = indices_pad.shape[0] + device = indices_pad.device d_probs = d_probs.contiguous() # d_routing_map ignored (bool tensor, no grad flow). @@ -444,18 +481,20 @@ def backward(ctx, d_probs, d_routing_map): # type: ignore[override] grid = (triton.cdiv(N, block_n),) _v4_router_post_bwd_kernel[grid]( d_probs, - indices_c, + indices_pad, scores_saved, weights_saved, denom_saved, d_logits_fp32, N, - E=E, - K=K, + 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 @@ -480,21 +519,6 @@ def is_triton_path_enabled() -> bool: return os.environ.get("PRIMUS_V4_ROUTER_TRITON", "1") != "0" -def is_triton_kernel_supported(logits: torch.Tensor, indices: torch.Tensor) -> bool: - """Return True iff inputs are CUDA + power-of-2 E + K.""" - if not logits.is_cuda or not indices.is_cuda: - return False - if logits.dim() != 2 or indices.dim() != 2: - return False - N, E = logits.shape - _, K = indices.shape - if E & (E - 1) != 0: - return False - if K & (K - 1) != 0: - return False - return True - - def v4_router_post_triton( logits: torch.Tensor, indices: torch.Tensor, @@ -508,7 +532,12 @@ def v4_router_post_triton( 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) @@ -516,5 +545,4 @@ def v4_router_post_triton( "V4RouterPostFn", "v4_router_post_triton", "is_triton_path_enabled", - "is_triton_kernel_supported", ] 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 index bde348bd2..883b233c9 100644 --- a/primus/backends/megatron/core/transformer/moe/v4_hash_router.py +++ b/primus/backends/megatron/core/transformer/moe/v4_hash_router.py @@ -56,9 +56,6 @@ 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_kernel_supported as _v4_router_triton_supported, -) from primus.backends.megatron.core.transformer.moe._triton.v4_router_post import ( is_triton_path_enabled as _v4_router_triton_enabled, ) @@ -204,11 +201,16 @@ def forward( 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 - 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())}" - ) + # 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)) @@ -218,7 +220,7 @@ def forward( # Plan-6 P39: route the post-logits chain through one fused # Triton kernel when PRIMUS_V4_ROUTER_TRITON=1 (default OFF). - if _v4_router_triton_enabled() and _v4_router_triton_supported(logits, indices): + if _v4_router_triton_enabled(): probs, routing_map = v4_router_post_triton( logits, indices, diff --git a/primus/backends/megatron/core/transformer/moe/v4_moe.py b/primus/backends/megatron/core/transformer/moe/v4_moe.py index 2a8bb795c..0748c83f5 100644 --- a/primus/backends/megatron/core/transformer/moe/v4_moe.py +++ b/primus/backends/megatron/core/transformer/moe/v4_moe.py @@ -86,6 +86,9 @@ 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, ) @@ -343,11 +346,15 @@ def _build_shared_expert_module(self, *, intermediate_size: int) -> nn.Module: shared_expert_spec, ModuleSpec ), "DeepSeek-V4 MoE requires shared_expert ModuleSpec in submodules." shared_expert_module = shared_expert_spec.module - assert shared_expert_module is SharedExpertMLP, "DeepSeek-V4 shared_expert must be SharedExpertMLP." + 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 must always run with clamped SwiGLU via SharedExpertMLP. + # 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 @@ -365,16 +372,24 @@ def _build_shared_expert_module(self, *, intermediate_size: int) -> nn.Module: 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( - shared_expert_spec, + 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 SharedExpertMLP: {exc}" + f"DeepSeek-V4 MoE shared expert build failed with PrimusSharedExpertMLP: {exc}" ) from exc def _build_token_dispatcher(self) -> MoETokenDispatcher: diff --git a/primus/backends/megatron/core/transformer/moe/v4_topk_router.py b/primus/backends/megatron/core/transformer/moe/v4_topk_router.py index 2ace392be..4d5b33252 100644 --- a/primus/backends/megatron/core/transformer/moe/v4_topk_router.py +++ b/primus/backends/megatron/core/transformer/moe/v4_topk_router.py @@ -65,9 +65,6 @@ 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_kernel_supported as _v4_router_triton_supported, -) from primus.backends.megatron.core.transformer.moe._triton.v4_router_post import ( is_triton_path_enabled as _v4_router_triton_enabled, ) @@ -131,7 +128,7 @@ def _compute_route( sel_score = scores_for_selection indices = sel_score.topk(topk, dim=-1).indices # [N, K] - if _v4_router_triton_enabled() and _v4_router_triton_supported(logits, indices): + 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 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 index cd84e7bf3..018e57a48 100644 --- a/primus/backends/megatron/core/transformer/v4_attention_kernels/__init__.py +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/__init__.py @@ -4,51 +4,216 @@ # See LICENSE for license information. ############################################################################### -""" -DeepSeek-V4 attention kernels (Plan-4). - -Plan-4 P24 lands two pure-function eager-Python references for V4 -attention — :func:`eager_v4_attention` (dense ``compress_ratio == 0`` / -HCA ``compress_ratio == 128``) and :func:`eager_v4_csa_attention` -(``compress_ratio == 4``). They are extracted from the math that -previously lived inline in :meth:`DeepseekV4Attention._attention_forward` -and :meth:`DeepseekV4Attention._csa_forward` so that: - -* there is exactly one definition of "the eager truth" that - ``DeepseekV4Attention``, the plan-4 Triton kernels (P25 / P26), and - the plan-4 unit-test harness all share; -* the existing checkpoint-reproduction baseline does not move (the - refactor is bit-identical at the existing call sites — see G22 - in ``tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p24_reference_op.py``); -* the reference op signatures match the future Triton kernel signatures - ``v4_attention`` (P25) and ``v4_csa_attention`` (P26) so the test - harness can plug reference ↔ candidate interchangeably. - -The Triton kernels themselves land in P25 / P26 alongside this package. +"""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_`` +* CSA (cr=4) entry: ``v4_csa_attention_`` + +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.reference import ( +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.v4_attention import ( - V4AttentionFn, - v4_attention, -) -from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention import ( +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_csa_attention, - v4_csa_attention_from_pool, + 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", - "v4_attention", - "V4AttentionFn", - "v4_csa_attention", + # triton v0 (deprecated gathered CSA) + "v4_csa_attention_v0", "V4CSAAttentionFn", - "v4_csa_attention_from_pool", + # 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/reference.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_eager/reference.py similarity index 99% rename from primus/backends/megatron/core/transformer/v4_attention_kernels/reference.py rename to primus/backends/megatron/core/transformer/v4_attention_kernels/_eager/reference.py index 3877f4f47..6fa6598d4 100644 --- a/primus/backends/megatron/core/transformer/v4_attention_kernels/reference.py +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_eager/reference.py @@ -15,7 +15,7 @@ * ``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`` / ``v4_csa_attention``) so the test harness can + (``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 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/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/__init__.py new file mode 100644 index 000000000..e69de29bb 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..f9ffcb6f3 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_bwd_flydsl_mqa.py @@ -0,0 +1,1270 @@ +"""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..16eb93c0e --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_fwd_flydsl_csa.py @@ -0,0 +1,219 @@ +"""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..70523c255 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_fwd_flydsl_mqa.py @@ -0,0 +1,143 @@ +"""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..373128656 --- /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,399 @@ +"""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..9f702010a --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_bwd_full_kernel.py @@ -0,0 +1,664 @@ +# SPDX-License-Identifier: Apache-2.0 +"""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..58362ad8f --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_bwd_dq_flydsl_kernel.py @@ -0,0 +1,592 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. + +"""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..4d74fafe9 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_fwd_v4_flydsl_kernel.py @@ -0,0 +1,650 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. + +"""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..3d38c10ab --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/__init__.py @@ -0,0 +1,37 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# 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..0626c846e --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/_dsa_bwd_gather.py @@ -0,0 +1,617 @@ +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..6315f6ff1 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/_dsa_bwd_preprocess.py @@ -0,0 +1,50 @@ +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..0529e0764 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_dkv_interm.py @@ -0,0 +1,229 @@ +""" +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..11b61b4e5 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_dq.py @@ -0,0 +1,502 @@ +""" +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..4be3b9a17 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_v4_gluon.py @@ -0,0 +1,166 @@ +""" +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..5cc778e19 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_fwd_v4_gluon.py @@ -0,0 +1,625 @@ +""" +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 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..dc2098b2a --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_dkv_interm_gluon.py @@ -0,0 +1,237 @@ +""" +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..78e6f58e1 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_dq_gluon.py @@ -0,0 +1,523 @@ +""" +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..e41e1a510 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_fwd_v4_gluon.py @@ -0,0 +1,722 @@ +""" +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 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..3ea978efc --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/aiter_lse_fwd.py @@ -0,0 +1,172 @@ +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..811a74428 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/aiter_mla_gluon.py @@ -0,0 +1,999 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +# 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..dc2098b2a --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_dkv_interm_gluon.py @@ -0,0 +1,237 @@ +""" +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..78e6f58e1 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_dq_gluon.py @@ -0,0 +1,523 @@ +""" +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..5c89b83e5 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_fwd_v4_gluon.py @@ -0,0 +1,733 @@ +""" +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 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 index 06a1ac8c5..97ad10bf1 100644 --- a/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/__init__.py +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/__init__.py @@ -33,7 +33,7 @@ - :func:`v4_csa_attention_bwd_tilelang` (P55 — cr=4 BWD) The dispatcher precedence (enforced inside the V4 attention -functional wrappers ``v4_attention`` / ``v4_csa_attention``) is: +functional wrappers ``v4_attention_v1`` / ``v4_csa_attention_v0``) is: .. code-block:: text @@ -291,7 +291,7 @@ def v4_csa_attention_bwd_tilelang(*args: Any, **kwargs: Any): # type: ignore[no # --------------------------------------------------------------------------- -# Dispatcher helpers — used by the v4_attention / v4_csa_attention wrappers +# Dispatcher helpers — used by the v4_attention_v1 / v4_csa_attention_v0 wrappers # --------------------------------------------------------------------------- _FALLBACK_WARNED: Set[str] = set() @@ -337,8 +337,8 @@ def _maybe_warn_fallback(kernel_name: str) -> None: def should_dispatch(kernel_name: str, enabled: bool = False) -> bool: - """Single dispatcher predicate used by ``v4_attention`` / - ``v4_csa_attention`` wrappers. + """Single dispatcher predicate used by ``v4_attention_v1`` / + ``v4_csa_attention_v0`` wrappers. Returns True iff all three conditions hold: 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 index 454686c92..10769790d 100644 --- 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 @@ -125,7 +125,7 @@ def v4_attention_tilelang( scale: Optional[float] = None, hca_local_seqlen: int = 0, ) -> torch.Tensor: - """Convenience wrapper matching the existing `v4_attention()` + """Convenience wrapper matching the existing `v4_attention_v1()` functional API but routing through the tilelang autograd Function.""" if scale is None: 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 index 696da06f3..937f8a70e 100644 --- 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 @@ -41,12 +41,12 @@ from typing import Optional, Tuple -import torch - 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.v4_attention_bwd import ( +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_bwd import ( _launch_v4_attention_bwd, ) 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 index 4e8fdafe2..7340b03c9 100644 --- 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 @@ -57,15 +57,15 @@ from typing import Optional, Tuple -import torch - # 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.v4_attention_fwd import ( +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_fwd import ( _launch_v4_attention_fwd, ) @@ -428,7 +428,7 @@ def v4_attention_fwd_tilelang( ) -> torch.Tensor: """Plan-8 P50 tilelang dense / SWA / sink FWD wrapper. - Drop-in replacement for the relevant subset of ``v4_attention`` + 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 diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/__init__.py deleted file mode 100644 index 92243e32c..000000000 --- a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/__init__.py +++ /dev/null @@ -1,29 +0,0 @@ -############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### - -"""Triton kernels backing :mod:`primus...transformer.v4_attention_kernels`. - -Submodules: - -* :mod:`v4_attention_fwd` — forward kernel for the dense - (``compress_ratio == 0``) and HCA (``compress_ratio == 128``) paths. -* :mod:`v4_attention_bwd` — backward kernel matching the FWD's saved - LSE. -* :mod:`v4_csa_attention_fwd` — forward kernel for the CSA - (``compress_ratio == 4``) fused local-SWA + per-query top-K + sink - path. -* :mod:`v4_csa_attention_bwd` — backward kernel for the CSA path - (``dq, dk_local, dv_local, dgathered, dsink``). -* :mod:`rope_interleaved_partial` — plan-6 P35 fused interleaved partial - RoPE FWD/BWD (replaces the 9-op eager chain in - :func:`primus.backends.megatron.core.transformer.dual_rope.apply_interleaved_partial_rope`). -* :mod:`sinkhorn` — plan-6 P36 fused Sinkhorn-Knopp FWD/BWD (replaces - the plan-5 P29 ``torch.compile`` fast path in - :func:`primus.backends.megatron.core.transformer.hyper_connection.sinkhorn_normalize`). - Runs the full alternating row/col normalize trajectory in registers - per row of the leading axis; BWD recomputes the trajectory and walks - the analytic VJP backward step by step. -""" 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/hc_glue.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_glue.py similarity index 100% rename from primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/hc_glue.py rename to primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_glue.py diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/indexer_score.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score.py similarity index 96% rename from primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/indexer_score.py rename to primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score.py index 3adf5a098..c82543d75 100644 --- a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/indexer_score.py +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score.py @@ -307,7 +307,17 @@ def _pick_block(s: int, p: int, h: int, hd: int) -> tuple[int, int]: if hd >= 256: block_s = 16 block_p = 32 - return min(block_s, triton.next_power_of_2(max(1, s))), min(block_p, triton.next_power_of_2(max(1, p))) + # 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 # --------------------------------------------------------------------------- diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/indexer_score_post.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score_post.py similarity index 100% rename from primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/indexer_score_post.py rename to primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score_post.py 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/rope_interleaved_partial.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rope_interleaved_partial.py similarity index 66% rename from primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/rope_interleaved_partial.py rename to primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rope_interleaved_partial.py index 8ae577a27..9aeaf2cd4 100644 --- a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/rope_interleaved_partial.py +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rope_interleaved_partial.py @@ -448,6 +448,241 @@ def backward(ctx, dout: torch.Tensor): # type: ignore[override] 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, +): + 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 + 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, +): + 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 + 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. + """ + + @staticmethod + def forward(ctx, x, pos_flat, inv_freq, rotary_dim): # 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), + ) + + ctx.save_for_backward(pos_c, inv_c) + ctx.rotary_dim = rotary_dim + ctx.head_dim = head_dim + ctx.H = H + ctx.N = N + 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), + ) + return dx, None, None, None + + +def apply_rope_from_positions( + x: torch.Tensor, + position_ids: torch.Tensor, + inv_freq: torch.Tensor, + *, + rotary_dim: int, +) -> 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. + """ + 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) + # Eager fallback: build cos/sin then apply. + freqs = position_ids.float().unsqueeze(-1) * inv_freq + return eager_apply_interleaved_partial_rope(x, freqs.cos(), freqs.sin(), rotary_dim=rotary_dim) + + # --------------------------------------------------------------------------- # Public Python entry points # --------------------------------------------------------------------------- @@ -535,7 +770,9 @@ def apply_rope_interleaved_partial( __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/sinkhorn.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/sinkhorn.py similarity index 100% rename from primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/sinkhorn.py rename to primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/sinkhorn.py 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/v4_attention.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention.py similarity index 90% rename from primus/backends/megatron/core/transformer/v4_attention_kernels/v4_attention.py rename to primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention.py index c68763f08..3acbc6493 100644 --- a/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_attention.py +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention.py @@ -8,7 +8,7 @@ Public API: -* :func:`v4_attention` — functional API matching +* :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 @@ -42,11 +42,14 @@ 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.v4_attention_bwd import ( +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.v4_attention_fwd import ( +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_fwd import ( _launch_v4_attention_fwd, ) @@ -79,7 +82,7 @@ def forward( # type: ignore[override] # dropout configuration raises rather than silently dropping # the kernel path. raise NotImplementedError( - "v4_attention does not implement in-kernel attention " + "v4_attention_v1 does not implement in-kernel attention " "dropout (V4 trains with attn_dropout=0). Got " f"attn_dropout={attn_dropout}, training={training}." ) @@ -175,7 +178,7 @@ def backward(ctx, grad_out: torch.Tensor): # type: ignore[override] # --------------------------------------------------------------------------- -def v4_attention( +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] @@ -188,6 +191,7 @@ def v4_attention( 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. @@ -258,6 +262,22 @@ def v4_attention( 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, @@ -274,5 +294,5 @@ def v4_attention( __all__ = [ "V4AttentionFn", - "v4_attention", + "v4_attention_v1", ] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/v4_attention_bwd.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_bwd.py similarity index 98% rename from primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/v4_attention_bwd.py rename to primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_bwd.py index 915be6354..ff8072ef9 100644 --- a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/v4_attention_bwd.py +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_bwd.py @@ -1454,10 +1454,10 @@ def _launch_v4_attention_bwd( ``dsink`` returned only when ``sink is not None`` (else ``None``). """ if not q.is_cuda: - raise ValueError("v4_attention BWD requires CUDA / HIP tensors.") + raise ValueError("v4_attention_v1 BWD requires CUDA / HIP tensors.") if dout.shape != out.shape or out.shape != q.shape: raise ValueError( - "v4_attention BWD shape mismatch: " + "v4_attention_v1 BWD shape mismatch: " f"out={tuple(out.shape)}, dout={tuple(dout.shape)}, q={tuple(q.shape)}" ) @@ -1689,9 +1689,20 @@ def _launch_v4_attention_bwd( # >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: - target = int(os.getenv("PRIMUS_V4_ATTN_BWD_DKV_HEAD_GROUPS", "1")) + # 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) @@ -1902,7 +1913,11 @@ def _launch_v4_attention_bwd( # 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). - pool_block_m = int(os.getenv("PRIMUS_V4_ATTN_BWD_POOL_BLOCK_M", str(BLOCK_M))) + # 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, diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/v4_attention_fwd.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_fwd.py similarity index 94% rename from primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/v4_attention_fwd.py rename to primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_fwd.py index 27ffd4987..4d1bf440e 100644 --- a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/v4_attention_fwd.py +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_fwd.py @@ -344,7 +344,7 @@ def _launch_v4_attention_fwd( """ if q.dim() != 4 or k.dim() != 4 or v.dim() != 4: raise ValueError( - "v4_attention forward expects q / k / v of rank 4 " + "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 @@ -352,15 +352,16 @@ def _launch_v4_attention_fwd( 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 shape mismatch: " f"q={tuple(q.shape)}, k={tuple(k.shape)}, v={tuple(v.shape)}" + "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 requires K_H ∈ {{1 (MQA), {HQ} (MHA)}}; got K_H={HK}.") + 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 requires CUDA / HIP tensors.") + raise ValueError("v4_attention_v1 requires CUDA / HIP tensors.") if q.dtype != k.dtype or q.dtype != v.dtype: raise ValueError( - "v4_attention requires q.dtype == k.dtype == v.dtype " f"(got {q.dtype} / {k.dtype} / {v.dtype})." + "v4_attention_v1 requires q.dtype == k.dtype == v.dtype " + f"(got {q.dtype} / {k.dtype} / {v.dtype})." ) has_sink = sink is not None @@ -414,11 +415,16 @@ def _launch_v4_attention_fwd( # 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`). - BLOCK_M = int(os.getenv("PRIMUS_V4_ATTN_FWD_BLOCK_M", "64")) - BLOCK_N = int(os.getenv("PRIMUS_V4_ATTN_FWD_BLOCK_N", "16")) - NUM_WARPS_FWD = int(os.getenv("PRIMUS_V4_ATTN_FWD_WARPS", "8")) - _default_stages_fwd = "1" if hca_local_seqlen else "2" - NUM_STAGES_FWD = int(os.getenv("PRIMUS_V4_ATTN_FWD_STAGES", _default_stages_fwd)) + # 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) 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/v4_csa_attention_bwd.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_bwd.py similarity index 86% rename from primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/v4_csa_attention_bwd.py rename to primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_bwd.py index 78eaa205c..030e35f76 100644 --- a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/v4_csa_attention_bwd.py +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_bwd.py @@ -56,7 +56,7 @@ Edge cases: * ``K_topk == 0`` — the wrapper short-circuits to the dense - :func:`v4_attention` BWD before reaching this kernel. + :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). @@ -71,7 +71,7 @@ import triton import triton.language as tl -from primus.backends.megatron.core.transformer.v4_attention_kernels._triton.v4_attention_bwd import ( +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, @@ -83,257 +83,6 @@ # --------------------------------------------------------------------------- -@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) - - @triton.jit def _v4_csa_attention_pool_bwd_kernel( Q, @@ -1414,164 +1163,6 @@ def _v4_csa_attention_pool_segreduce_kernel( # --------------------------------------------------------------------------- -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 BWD requires CUDA / HIP tensors.") - if dout.shape != out.shape or out.shape != q.shape: - raise ValueError( - "v4_csa_attention 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 - - def _launch_v4_csa_attention_pool_bwd( q: torch.Tensor, # [B, H, Sq, D] k_local: torch.Tensor, # [B, H, Sq, D] @@ -1588,17 +1179,19 @@ def _launch_v4_csa_attention_pool_bwd( ) -> 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 pool BWD requires CUDA / HIP tensors.") + 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 pool BWD shape mismatch: " + "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 pool BWD expects pool rank 3 [B, P, D], got {tuple(pool.shape)}.") + 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 pool BWD expects topk_idxs rank 3 [B, Sq, K], got {tuple(topk_idxs.shape)}." + 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 @@ -1606,11 +1199,11 @@ def _launch_v4_csa_attention_pool_bwd( Bt, Sqt, K_topk = topk_idxs.shape if Bp != B or Dp != D or Bt != B or Sqt != Sq: raise ValueError( - "v4_csa_attention pool BWD shape mismatch: " + "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 topk_idxs must be int32/int64, got {topk_idxs.dtype}.") + raise ValueError(f"v4_csa_attention_v0 topk_idxs must be int32/int64, got {topk_idxs.dtype}.") has_sink = sink is not None diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/v4_csa_attention_fwd.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_fwd.py similarity index 72% rename from primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/v4_csa_attention_fwd.py rename to primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_fwd.py index 19d0f6d56..5ee75bbbe 100644 --- a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton/v4_csa_attention_fwd.py +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_fwd.py @@ -38,7 +38,7 @@ Edge cases handled: * ``K_topk == 0`` — wrapper short-circuits to the dense - :func:`v4_attention` kernel before reaching this file. + :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. @@ -60,7 +60,7 @@ import triton import triton.language as tl -from primus.backends.megatron.core.transformer.v4_attention_kernels._triton.v4_attention_fwd import ( +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_fwd import ( _launch_v4_attention_fwd, ) @@ -69,217 +69,6 @@ # --------------------------------------------------------------------------- -@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) - - @triton.jit def _v4_csa_attention_pool_fwd_kernel( Q, @@ -443,138 +232,6 @@ def _v4_csa_attention_pool_fwd_kernel( # --------------------------------------------------------------------------- -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 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 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 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 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 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 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 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 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 - - def _launch_v4_csa_attention_pool_fwd( q: torch.Tensor, # [B, H, Sq, D] k_local: torch.Tensor, # [B, H, Sq, D] @@ -589,46 +246,46 @@ def _launch_v4_csa_attention_pool_fwd( """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 pool forward expects q / k_local / v_local of rank 4 " + "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 pool forward expects pool of rank 3 [B, P, D]; " + 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 pool forward expects topk_idxs of rank 3 [B, Sq, K]; " + 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 pool path requires k_local.shape == v_local.shape == q.shape " + "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 pool shape mismatch: expected " + "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 topk_idxs shape mismatch: expected " + "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 topk_idxs must be int32/int64, got {topk_idxs.dtype}.") + 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 requires CUDA / HIP tensors.") + 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 pool path requires q.dtype == k_local.dtype == " + "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})." ) @@ -1217,9 +874,15 @@ def _launch_v4_csa_attention_pool_fwd_split( # 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 = 64 if HQ >= 64 else triton.next_power_of_2(HQ) - if BLOCK_H > HQ: - BLOCK_H = max(triton.next_power_of_2(HQ), 16) + # 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) 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..5381d3d46 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_turbo_flydsl/__init__.py @@ -0,0 +1,41 @@ +############################################################################### +# 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 v2 attention +(``primus_turbo.flydsl.attention.kernels.sparse_mla_v2``) from the Primus-Turbo +``dev/kyle/flydsl_attn_deepseekv4`` line. 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)`` + +``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 + +from primus_turbo.flydsl.attention.kernels.sparse_mla_v2 import ( + sparse_mla_bwd_v4_flydsl, + sparse_mla_fwd_v4_flydsl, +) + +# 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_v4_flydsl +sparse_mla_bwd_v4_turbo_flydsl = sparse_mla_bwd_v4_flydsl + +__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.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention.py deleted file mode 100644 index 5242e50ae..000000000 --- a/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention.py +++ /dev/null @@ -1,412 +0,0 @@ -############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### - -"""V4 CSA attention Triton autograd entry point (plan-4 P26). - -Public API: - -* :func:`v4_csa_attention` — functional API matching - :func:`eager_v4_csa_attention`'s signature; routes through - :class:`V4CSAAttentionFn` so autograd works. -* :class:`V4CSAAttentionFn` — :class:`torch.autograd.Function` wrapping - the Triton FWD + Triton BWD (re-materialises the joint softmax row - from the saved LSE; emits ``dq, dk_local, dv_local, dgathered, - dsink``). - -Dispatch contract (consumed by ``DeepseekV4Attention._csa_forward``): - -* ``compress_ratio == 4`` only — the dense / HCA paths use the separate - :func:`v4_attention` kernel. -* The wrapper takes the already-gathered ``[B, Sq, K_topk, head_dim]`` - tensor; the per-query top-K gather lives in ``_csa_forward`` outside - the kernel (matches plan-4 design — see Phase 26 design notes). -* When ``K_topk == 0`` (degenerate Indexer state), the wrapper - short-circuits to the dense :func:`v4_attention` kernel; the local - SWA + sink path is mathematically identical to CSA's local-only - branch in that limit. -* When ``attn_dropout > 0`` and ``training=True`` the wrapper raises - :class:`NotImplementedError` (V4 trains with ``attn_dropout=0``; the - kernel does not implement in-kernel dropout). This matches the dense - :func:`v4_attention` contract. - -dtype contract (matches :func:`eager_v4_csa_attention`): - -* Q / K_local / V_local / gathered matmuls run in input dtype with fp32 - accumulator (computed via fp32 upcast inside the kernel — see - ``v4_csa_attention_fwd`` docstring). -* The online-softmax accumulator (``m_running / l_running / acc``) is - fp32 — the *only* fp32 step. -* Output is in ``v_local.dtype``. -* Saved LSE is fp32 (BWD re-materialises the joint ``P`` row from it). - -Gradient routing: - -The wrapper's BWD returns ``dgathered`` as the gradient w.r.t. the -*post-gather* tensor — the caller (``DeepseekV4Attention._csa_forward`` -replacement) is responsible for the scatter-add back into the -compressed-pool gradient via the autograd of ``torch.gather`` / -``torch.expand``. This matches how PyTorch's eager autograd already -works (see the eager reference in -:func:`eager_v4_csa_attention`'s docstring). -""" - -from __future__ import annotations - -from typing import Optional - -import torch - -from primus.backends.megatron.core.transformer.v4_attention_kernels import _tilelang -from primus.backends.megatron.core.transformer.v4_attention_kernels._triton.v4_csa_attention_bwd import ( - _launch_v4_csa_attention_bwd, - _launch_v4_csa_attention_pool_bwd, -) -from primus.backends.megatron.core.transformer.v4_attention_kernels._triton.v4_csa_attention_fwd import ( - _launch_v4_csa_attention_fwd, - _launch_v4_csa_attention_pool_fwd, -) -from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_attention import ( - v4_attention, -) - -# --------------------------------------------------------------------------- -# Autograd Function -# --------------------------------------------------------------------------- - - -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 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 - - -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_from_pool 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 - - -# --------------------------------------------------------------------------- -# Functional API -# --------------------------------------------------------------------------- - - -def 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_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, -) -> 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` 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's dense+SWA+sink - # path bit-identically when K_topk == 0 (the joint softmax - # collapses to the local-only softmax). - return v4_attention( - 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, - ) - return V4CSAAttentionFn.apply( - q, - k_local, - v_local, - gathered, - sparse_mask, - sink, - swa_window, - attn_dropout, - training, - scale, - ) - - -def v4_csa_attention_from_pool( - 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, -) -> 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( - 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__ = [ - "V4CSAAttentionFn", - "V4CSAPoolAttentionFn", - "v4_csa_attention", - "v4_csa_attention_from_pool", -] 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/patches/deepseek_v4_get_batch_patches.py b/primus/backends/megatron/patches/deepseek_v4_get_batch_patches.py index ff4936a99..66269e6c5 100644 --- a/primus/backends/megatron/patches/deepseek_v4_get_batch_patches.py +++ b/primus/backends/megatron/patches/deepseek_v4_get_batch_patches.py @@ -80,7 +80,6 @@ from primus.core.patches import PatchContext, get_args, register_patch from primus.modules.module_utils import log_rank_0 - # --------------------------------------------------------------------------- # Module-local pre-broadcast cache. # @@ -240,10 +239,7 @@ def patch_v4_pp_token_pre_broadcast(ctx: PatchContext): 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" - ) + 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. diff --git a/primus/backends/megatron/patches/deepseek_v4_pp_shape_patches.py b/primus/backends/megatron/patches/deepseek_v4_pp_shape_patches.py index f06efef19..ca6ae9f81 100644 --- a/primus/backends/megatron/patches/deepseek_v4_pp_shape_patches.py +++ b/primus/backends/megatron/patches/deepseek_v4_pp_shape_patches.py @@ -132,14 +132,9 @@ def patch_v4_pp_tensor_shape(ctx: PatchContext): # 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" - ) + 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 - ) + 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} " @@ -153,12 +148,11 @@ def patch_v4_pp_tensor_shape(ctx: PatchContext): 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" + "[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) + 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 " 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..852e80462 --- /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.modules.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/mla_patches.py b/primus/backends/megatron/patches/mla_patches.py index 05a3b1d60..547482c9a 100644 --- a/primus/backends/megatron/patches/mla_patches.py +++ b/primus/backends/megatron/patches/mla_patches.py @@ -21,9 +21,25 @@ phase="before_train", description=( "Monkey patch MLA attention to use PrimusMLASelfAttention " - "when use_turbo_parallel_linear is enabled." + "when use_turbo_parallel_linear 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_parallel_linear", 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_parallel_linear", False), ) def patch_mla_attention(ctx: PatchContext): """ 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..7b7be0f8a --- /dev/null +++ b/primus/backends/megatron/patches/moe_alltoall_dtoh_patches.py @@ -0,0 +1,107 @@ +############################################################################### +# 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.modules.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_mlp`` + is the flag users actually set; ``use_turbo_grouped_gemm`` is the alias + auto-enabled by Sync-Free MoE stage >= 2. + """ + try: + args = get_args(ctx) + except Exception: + return False + return bool( + getattr(args, "use_turbo_grouped_mlp", False) or 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_mlp) 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/turbo/gpt_output_layer_patches.py b/primus/backends/megatron/patches/turbo/gpt_output_layer_patches.py index 6a9ed9a60..249866b98 100644 --- a/primus/backends/megatron/patches/turbo/gpt_output_layer_patches.py +++ b/primus/backends/megatron/patches/turbo/gpt_output_layer_patches.py @@ -12,8 +12,7 @@ import importlib.util -from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.patches import PatchContext, get_args def _is_turbo_parallel_linear_enabled(ctx: PatchContext) -> bool: @@ -36,33 +35,3 @@ def _is_turbo_parallel_linear_enabled(ctx: PatchContext) -> bool: use_turbo_parallel_linear = bool(getattr(args, "use_turbo_parallel_linear", False)) return tp_size == 1 and enable_primus_turbo and use_turbo_parallel_linear - - -@register_patch( - "megatron.turbo.gpt_output_layer", - backend="megatron", - phase="before_train", - description="Replace GPT ColumnParallelLinear with PrimusTurbo implementation", - condition=_is_turbo_parallel_linear_enabled, -) -def patch_gpt_output_layer(ctx: PatchContext): - """ - Patch GPT output layer to use PrimusTurbo ColumnParallelLinear. - - This replaces the standard tensor_parallel.ColumnParallelLinear with - PrimusTurboColumnParallelLinearTorch in gpt_model. - """ - from megatron.core.models.gpt import gpt_model - - from primus.backends.megatron.core.extensions.primus_turbo import ( - PrimusTurboColumnParallelLinearTorch, - ) - - log_rank_0("[Patch:megatron.turbo.gpt_output_layer] Patching GPT output layer...") - - gpt_model.tensor_parallel.ColumnParallelLinear = PrimusTurboColumnParallelLinearTorch - log_rank_0( - "[Patch:megatron.turbo.gpt_output_layer] Patched " - f"megatron.core.models.gpt.gpt_model.tensor_parallel.ColumnParallelLinear " - f"-> {PrimusTurboColumnParallelLinearTorch.__name__}" - ) diff --git a/primus/backends/transformer_engine/pytorch/permutation.py b/primus/backends/transformer_engine/pytorch/permutation.py index 331f1ea0e..29d2ca139 100644 --- a/primus/backends/transformer_engine/pytorch/permutation.py +++ b/primus/backends/transformer_engine/pytorch/permutation.py @@ -585,6 +585,7 @@ def moe_unpermute( restore_shape: Optional[torch.Size] = None, map_type: str = "mask", probs: Optional[torch.Tensor] = None, + pad_offsets: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Unpermute a tensor with permuted tokens, and optionally merge the tokens with their @@ -608,6 +609,11 @@ def moe_unpermute( Options are: 'mask', 'index'. probs: torch.Tensor, default = None Renamed to merging_probs. Keep for backward compatibility. + pad_offsets: torch.Tensor, default = None + Accepted for API compatibility with upstream TE >= 2.12.0 + ``moe_unpermute`` (Megatron passes it when ``is_te_min_version("2.12.0")``). + This ROCm/Triton mask-map path does not pad during permutation, so the + argument is ignored. """ if probs is not None: if merging_probs is not None: diff --git a/primus/configs/models/megatron/deepseek_v4_base.yaml b/primus/configs/models/megatron/deepseek_v4_base.yaml index 7b65215ae..bd874d588 100644 --- a/primus/configs/models/megatron/deepseek_v4_base.yaml +++ b/primus/configs/models/megatron/deepseek_v4_base.yaml @@ -46,32 +46,21 @@ index_n_heads: 64 attn_sliding_window: 128 attn_sink: true -# Plan-4 P25 / P26: in-tree Primus Triton kernels for V4 attention. -# Defaults to false so existing checkpoints / smokes are unaffected; -# the run script flips them via PRIMUS_USE_V4_TRITON_ATTENTION / -# PRIMUS_USE_V4_TRITON_CSA_ATTENTION env vars (see run_deepseek_v4.sh). -# Precedence in DeepseekV4Attention.forward: -# use_turbo_attention > use_v4_tilelang_attention > use_v4_triton_attention > eager (cr ∈ {0, 128}) -# use_v4_tilelang_csa_attention > use_v4_triton_csa_attention > eager (cr == 4) -use_v4_triton_attention: false -use_v4_triton_csa_attention: false +# 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 -# Plan-8 tilelang attention kernels (OPTIONAL dependency, default OFF). -# Tilelang is NOT bundled in the default Primus container; enabling -# these flags without a tilelang install logs a one-time rank-0 -# warning and falls back to the plan-4 / plan-5 Triton path -- training -# continues without error. Surface via PRIMUS_USE_V4_TILELANG_ATTENTION -# / PRIMUS_USE_V4_TILELANG_CSA_ATTENTION env vars in run_deepseek_v4.sh. -# Replaces the legacy PRIMUS_V4_TILELANG_ATTN env knob; that env var -# is no longer consulted. -use_v4_tilelang_attention: false -use_v4_tilelang_csa_attention: false - # 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 @@ -121,6 +110,7 @@ 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 diff --git a/primus/configs/modules/megatron/primus_megatron_module.yaml b/primus/configs/modules/megatron/primus_megatron_module.yaml index 0ec3a22b0..a8cd951c1 100644 --- a/primus/configs/modules/megatron/primus_megatron_module.yaml +++ b/primus/configs/modules/megatron/primus_megatron_module.yaml @@ -18,7 +18,7 @@ use_rocm_mem_info_iters: [1,2] 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 diff --git a/primus/configs/modules/megatron/primus_turbo.yaml b/primus/configs/modules/megatron/primus_turbo.yaml index f0a90ae23..a4547d5b8 100644 --- a/primus/configs/modules/megatron/primus_turbo.yaml +++ b/primus/configs/modules/megatron/primus_turbo.yaml @@ -22,8 +22,6 @@ use_turbo_parallel_linear: false # ===== Grouped MLP ===== # operator switch use_turbo_grouped_mlp: false -# temporary switch: disable Turbo GroupedMLP low-precision grouped_gemm path -disable_turbo_grouped_mlp_low_precision: false # fused activation_with_probs to reduce redundant computation use_turbo_fused_act_with_probs: false 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 cbe3526db..272b62535 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 6a0ca72f1..cb9a1b8c8 100644 --- a/primus/core/projection/module_profilers/transformer_layer.py +++ b/primus/core/projection/module_profilers/transformer_layer.py @@ -19,7 +19,7 @@ from .moe_mlp import MoEMLPProfiler from .residual_add import ResidualAddProfiler from .router import RouterProfiler -from .utils import 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 @@ -315,11 +315,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 @@ -443,11 +459,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 diff --git a/primus/core/projection/module_profilers/utils.py b/primus/core/projection/module_profilers/utils.py index 111c03bff..9ec612351 100644 --- a/primus/core/projection/module_profilers/utils.py +++ b/primus/core/projection/module_profilers/utils.py @@ -55,11 +55,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. @@ -110,6 +152,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! @@ -126,7 +169,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,) @@ -160,7 +203,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) @@ -184,7 +227,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 --- @@ -221,6 +264,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. @@ -273,6 +317,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()} fp8_context = _get_fp8_context_for_benchmark(transformer_config) # ========================================================================= @@ -284,7 +329,7 @@ def create_input(spec): with fp8_context: for _ in range(num_warmup): - outputs = moe_module(*inputs) + outputs = moe_module(*inputs, **kwargs) if not isinstance(outputs, (tuple, list)): outputs = (outputs,) @@ -316,7 +361,7 @@ def create_input(spec): with fp8_context: for _ in range(num_iterations): - outputs = moe_module(*inputs) + outputs = moe_module(*inputs, **kwargs) if device.type == "cuda": torch.cuda.synchronize(device) @@ -372,7 +417,7 @@ def timed_combine(*args, **kwargs): forward_end = torch.cuda.Event(enable_timing=True) forward_start.record() - outputs = moe_module(*inputs) + outputs = moe_module(*inputs, **kwargs) forward_end.record() # --- Backward pass --- diff --git a/primus/core/projection/performance_projection/projection.py b/primus/core/projection/performance_projection/projection.py index d78ee0af7..d069911a3 100644 --- a/primus/core/projection/performance_projection/projection.py +++ b/primus/core/projection/performance_projection/projection.py @@ -521,10 +521,22 @@ def _limit_layers_for_projection(module_config): original_moe_layout = getattr(module_config, "moe_layer_freq", None) dense_layers_present = _has_dense_layers(original_moe_layout) # Use 1 layer for fast profiling - results are extrapolated to full model - # Increase to 2-4 for better accuracy if needed - max_layers = 1 + # 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: diff --git a/primus/modules/module_utils.py b/primus/modules/module_utils.py index 52503d985..b6f171895 100644 --- a/primus/modules/module_utils.py +++ b/primus/modules/module_utils.py @@ -86,7 +86,7 @@ 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(msg="", *args, **kwargs): log_func = logger.debug_with_caller caller = inspect.stack()[1] diff --git a/primus/modules/trainer/megatron/trainer.py b/primus/modules/trainer/megatron/trainer.py index 312df71ec..3a898a788 100644 --- a/primus/modules/trainer/megatron/trainer.py +++ b/primus/modules/trainer/megatron/trainer.py @@ -512,14 +512,19 @@ def update_primus_config( # 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": + if model_type != "gpt": 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": + # Pass model_type through for any non-gpt model (mamba, deepseek_v4, ...) + # so the correct builder is selected. Without this, deepseek_v4 falls + # back to the stock GPT builder (SelfAttention instead of + # DeepseekV4Attention) — which is what made the projection benchmark + # stock attention and ignore use_v4_triton_attention / compress_ratios. + if model_type != "gpt": 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}") @@ -610,8 +615,6 @@ def train_valid_test_datasets_provider_func(train_val_test_num_samples, vp_stage args.do_valid, args.do_test, args.dataloader_type, - args.retro_project_dir, - args.retro_cyclic_train_iters, ) # Print setup timing. @@ -916,9 +919,6 @@ def setup_model_and_optimizer( optimizer = get_megatron_optimizer( config, model, - no_wd_decay_cond, - scale_lr_cond, - lr_mult, use_gloo_process_groups=args.enable_gloo_process_groups, ) else: @@ -931,9 +931,6 @@ def setup_model_and_optimizer( optimizer = get_megatron_muon_optimizer( config, model, - no_wd_decay_cond, - scale_lr_cond, - lr_mult, use_gloo_process_groups=args.enable_gloo_process_groups, layer_wise_distributed_optimizer="dist" in config.optimizer, ) diff --git a/primus/modules/trainer/megatron/utils.py b/primus/modules/trainer/megatron/utils.py index 6bcc96e77..739038779 100644 --- a/primus/modules/trainer/megatron/utils.py +++ b/primus/modules/trainer/megatron/utils.py @@ -422,23 +422,31 @@ def dump_pp_data(args, num_mbs, pp_data_dir): json.dump(config_dict, f, indent=2) -def _get_sync_free_moe_options(stage: int) -> dict: +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}, + 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_mlp": 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_mlp": 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, }, } @@ -499,7 +507,7 @@ def validate_args_on_rocm(args): raise ValueError( "Sync-Free MoE stage 2 or 3 require PrimusTurboGroupedMLP, please set `use_turbo_grouped_mlp=True`" ) - options = _get_sync_free_moe_options(args.turbo_sync_free_moe_stage) + 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) ==========" ) diff --git a/rccl_avg_workaround/.gitignore b/rccl_avg_workaround/.gitignore new file mode 100644 index 000000000..7a60b85e1 --- /dev/null +++ b/rccl_avg_workaround/.gitignore @@ -0,0 +1,2 @@ +__pycache__/ +*.pyc diff --git a/rccl_avg_workaround/sitecustomize.py b/rccl_avg_workaround/sitecustomize.py new file mode 100644 index 000000000..bf10ac87e --- /dev/null +++ b/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/run_deepseek_v4.sh b/run_deepseek_v4.sh index d605bbe3f..93492074f 100755 --- a/run_deepseek_v4.sh +++ b/run_deepseek_v4.sh @@ -26,13 +26,16 @@ export TRAIN_ITERS=${TRAIN_ITERS:-20} export DOCKER_IMAGE=${DOCKER_IMAGE:-"tasimage/primus:pr-715-ainic"} export SLURM_PARTITION=Compute-DCPT -export SLURM_NODELIST="${SLURM_NODELIST:-smci355-ccs-aus-n01-21,smci355-ccs-aus-n01-33,smci355-ccs-aus-n02-25,smci355-ccs-aus-n02-33,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}" +export SLURM_NODELIST="${SLURM_NODELIST:-smci355-ccs-aus-n01-21,smci355-ccs-aus-n01-33,smci355-ccs-aus-n02-21,smci355-ccs-aus-n02-25,smci355-ccs-aus-n02-29,smci355-ccs-aus-n02-33,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}" 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" -export GLOO_SOCKET_IFNAME=fenic -export NCCL_SOCKET_IFNAME=fenic +# "fenic" is the cluster NIC; honor inherited values so a non-cluster / direct +# run (e.g. local single-node smoke) can point these at a real interface +# (e.g. lo) — otherwise gloo aborts with "Unable to find address for: fenic". +export GLOO_SOCKET_IFNAME=${GLOO_SOCKET_IFNAME:-fenic} +export NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME:-fenic} 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} @@ -68,6 +71,10 @@ if [ "$USE_TURBO_ATTENTION" = "True" ] || [ "${USE_TURBO_DEEPEP:-False}" = "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 @@ -143,31 +150,21 @@ if [ "$OPTIMIZER" = "muon" ] || [ "$OPTIMIZER" = "dist_muon" ]; then ) fi -# Plan-4 P25 / P26: in-tree Primus Triton kernels for V4 attention. -# Precedence in DeepseekV4Attention.forward: -# use_turbo_attention > use_v4_tilelang_attention > use_v4_triton_attention > eager (cr ∈ {0, 128}) -# use_v4_tilelang_csa_attention > use_v4_triton_csa_attention > eager (cr == 4) -# These are V4-only; they have no effect on other model types. -export USE_V4_TRITON_ATTENTION=${USE_V4_TRITON_ATTENTION:-True} -export USE_V4_TRITON_CSA_ATTENTION=${USE_V4_TRITON_CSA_ATTENTION:-True} +# 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:-triton_v2} +export USE_V4_CSA_ATTENTION_BACKEND=${USE_V4_CSA_ATTENTION_BACKEND:-triton_v2} # 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} -# Plan-8 tilelang attention kernels (OPTIONAL — default OFF). -# Tilelang is NOT bundled in the default Primus container, so we leave -# both flags off here. When the container has tilelang installed at -# the plan-8 pin AND the P50..P55 kernels are registered, set these -# to True (e.g. in a sweep / experiment script) to route the dense / -# HCA / CSA paths through tilelang. Otherwise the dispatcher prints -# a single rank-0 warning and falls back to the Triton path -- training -# continues without error. Replaces the legacy PRIMUS_V4_TILELANG_ATTN -# env knob (no longer consulted). -export USE_V4_TILELANG_ATTENTION=${USE_V4_TILELANG_ATTENTION:-False} -export USE_V4_TILELANG_CSA_ATTENTION=${USE_V4_TILELANG_CSA_ATTENTION:-False} - # 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 @@ -182,13 +179,16 @@ export USE_V4_COMPILED_SINKHORN=${USE_V4_COMPILED_SINKHORN:-False} # 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 { [ "$USE_V4_TRITON_ATTENTION" = "True" ] || [ "$USE_V4_TRITON_CSA_ATTENTION" = "True" ]; } && [ "${PRIMUS_TP:-1}" -gt 1 ]; then +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 - export FP8=${FP8:-hybrid} - export FP8_RECIPE=${FP8_RECIPE:-delayed} + # 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) ---- @@ -234,10 +234,20 @@ mkdir -p "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" export PRIMUS_EXIT_FAST=1 -./primus-cli slurm -N "$NNODES" \ - ${SLURM_PARTITION:+--partition="${SLURM_PARTITION}"} \ - ${SLURM_NODELIST:+--nodelist="${SLURM_NODELIST}"} \ - -- --image "${DOCKER_IMAGE}" --clean -- --numa \ +# 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) +else + LAUNCHER_ARGS=(slurm -N "$NNODES") + [ -n "${SLURM_PARTITION:-}" ] && LAUNCHER_ARGS+=(--partition="${SLURM_PARTITION}") + [ -n "${SLURM_NODELIST:-}" ] && LAUNCHER_ARGS+=(--nodelist="${SLURM_NODELIST}") + LAUNCHER_ARGS+=(-- --image "${DOCKER_IMAGE}" --clean -- --numa) +fi + +./primus-cli "${LAUNCHER_ARGS[@]}" \ -- train pretrain --config "$EXP" \ --manual_gc True \ --manual_gc_interval 100 \ @@ -269,11 +279,9 @@ export PRIMUS_EXIT_FAST=1 --mock_data True \ --enable_primus_turbo "$ENABLE_PRIMUS_TURBO" \ --use_turbo_attention "$USE_TURBO_ATTENTION" \ - --use_v4_triton_attention "$USE_V4_TRITON_ATTENTION" \ - --use_v4_triton_csa_attention "$USE_V4_TRITON_CSA_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_tilelang_attention "$USE_V4_TILELANG_ATTENTION" \ - --use_v4_tilelang_csa_attention "$USE_V4_TILELANG_CSA_ATTENTION" \ --use_v4_compiled_sinkhorn "$USE_V4_COMPILED_SINKHORN" \ --use_turbo_deepep "$USE_TURBO_DEEPEP" \ "${TURBO_DEEPEP_CLI_ARGS[@]}" \ @@ -295,4 +303,5 @@ export PRIMUS_EXIT_FAST=1 --use_pytorch_profiler "$PROFILE" \ --profile_step_end 7 \ --profile_step_start 6 \ + --bias_swiglu_fusion "$PRIMUS_BIAS_SWIGLU_FUSION" \ 2>&1 | tee "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME/log_node_${NODE_RANK:-0}.txt" diff --git a/run_deepseek_v4_flash.sh b/run_deepseek_v4_flash.sh index 8d2a3f035..5c5a0daa9 100644 --- a/run_deepseek_v4_flash.sh +++ b/run_deepseek_v4_flash.sh @@ -10,30 +10,30 @@ 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=8 -# export PRIMUS_TP=${PRIMUS_TP:-1} -# export PRIMUS_PP=${PRIMUS_PP:-8} -# export PRIMUS_EP=${PRIMUS_EP:-8} -# export PRIMUS_RECOMPUTE_LAYERS=0 -# if [ "$MTP_NUM_LAYERS" -eq 1 ]; then -# export PRIMUS_PP_LAYOUT='"Et*4|t*5|(t*6|)*5,t*4mL"' -# else -# export PRIMUS_PP_LAYOUT='"Et*4|t*5|(t*6|)*5,t*4L"' -# fi +export NNODES=${NNODES:-8} -export NNODES=4 -export PRIMUS_TP=${PRIMUS_TP:-1} -export PRIMUS_PP=${PRIMUS_PP:-4} -export PRIMUS_EP=${PRIMUS_EP:-8} -export PRIMUS_RECOMPUTE_LAYERS=3 -if [ "$MTP_NUM_LAYERS" -eq 1 ]; then - export PRIMUS_PP_LAYOUT='"Et*10|t*11|t*11|t*11mL"' -else - export PRIMUS_PP_LAYOUT='"Et*10|t*11|t*11|t*11L"' +if [ "$NNODES" -eq 8 ]; then + export PRIMUS_TP=${PRIMUS_TP:-1} + export PRIMUS_PP=${PRIMUS_PP:-8} + export PRIMUS_EP=${PRIMUS_EP:-8} + export PRIMUS_RECOMPUTE_LAYERS=0 + if [ "$MTP_NUM_LAYERS" -eq 1 ]; then + export PRIMUS_PP_LAYOUT='"Et*4|t*5|(t*6|)*5,t*4mL"' + else + export PRIMUS_PP_LAYOUT='"Et*4|t*5|(t*6|)*5,t*4L"' + fi +elif [ "$NNODES" -eq 4 ]; then + export PRIMUS_TP=${PRIMUS_TP:-1} + export PRIMUS_PP=${PRIMUS_PP:-4} + export PRIMUS_EP=${PRIMUS_EP:-8} + export PRIMUS_RECOMPUTE_LAYERS=3 + if [ "$MTP_NUM_LAYERS" -eq 1 ]; then + export PRIMUS_PP_LAYOUT='"Et*10|t*11|t*11|t*11mL"' + else + export PRIMUS_PP_LAYOUT='"Et*10|t*11|t*11|t*11L"' + fi fi -export SLURM_NODELIST=${SLURM_NODELIST:-"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"} - export MBS=${MBS:-1} export GBS=${GBS:-$((64 * NNODES * MBS))} export TRAIN_ITERS=${TRAIN_ITERS:-10} @@ -42,8 +42,8 @@ export PRIMUS_SEQ_LENGTH=${PRIMUS_SEQ_LENGTH:-4096} export PRIMUS_MAX_POSITION_EMBEDDINGS=${PRIMUS_MAX_POSITION_EMBEDDINGS:-${PRIMUS_SEQ_LENGTH}} export USE_V4_FP8_INDEXER=${USE_V4_FP8_INDEXER:-True} -export USE_V4_TRITON_ATTENTION=${USE_V4_TRITON_ATTENTION:-True} -export USE_V4_TRITON_CSA_ATTENTION=${USE_V4_TRITON_CSA_ATTENTION:-True} +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} export USE_V4_COMPILED_SINKHORN=${USE_V4_COMPILED_SINKHORN:-True} diff --git a/run_deepseek_v4_flash_proxy.sh b/run_deepseek_v4_flash_proxy.sh index 91ad20d95..1552fd6ff 100755 --- a/run_deepseek_v4_flash_proxy.sh +++ b/run_deepseek_v4_flash_proxy.sh @@ -23,8 +23,8 @@ # command line if OOM) # # Plan-5 perf knobs default ON: -# - USE_V4_TRITON_ATTENTION (cr ∈ {0, 128} -> Primus Triton kernel) -# - USE_V4_TRITON_CSA_ATTENTION (cr == 4 -> Primus Triton CSA) +# - 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, @@ -128,8 +128,8 @@ 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_TRITON_ATTENTION=${USE_V4_TRITON_ATTENTION:-True} -export USE_V4_TRITON_CSA_ATTENTION=${USE_V4_TRITON_CSA_ATTENTION:-True} +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 @@ -242,6 +242,17 @@ export PRIMUS_INDEXER_TRITON_FULL=${PRIMUS_INDEXER_TRITON_FULL:-0} # 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 @@ -255,6 +266,14 @@ export PROFILE=${PROFILE:-False} # 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)" diff --git a/run_deepseek_v4_pro_muon.sh b/run_deepseek_v4_pro_muon.sh index 438e43fa3..86257a16b 100755 --- a/run_deepseek_v4_pro_muon.sh +++ b/run_deepseek_v4_pro_muon.sh @@ -39,7 +39,13 @@ # 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 @@ -52,7 +58,7 @@ 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-BF16-pretrain.yaml} +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} @@ -120,6 +126,13 @@ export MUON_MOMENTUM=${MUON_MOMENTUM:-0.95} # paper momentum 0.95 # 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): @@ -141,13 +154,111 @@ 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_TRITON_ATTENTION=${USE_V4_TRITON_ATTENTION:-True} -export USE_V4_TRITON_CSA_ATTENTION=${USE_V4_TRITON_CSA_ATTENTION:-True} -export USE_V4_TILELANG_ATTENTION=${USE_V4_TILELANG_ATTENTION:-False} -export USE_V4_TILELANG_CSA_ATTENTION=${USE_V4_TILELANG_CSA_ATTENTION:-False} +# 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_mlp=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} @@ -187,7 +298,7 @@ mkdir -p "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" --manual_gc_interval 100 \ --num_layers "$PRIMUS_TOTAL_LAYERS" \ --train_iters "$TRAIN_ITERS" \ - --lr_warmup_iters 0 \ + --lr_warmup_iters "${LR_WARMUP_ITERS:-0}" \ --lr_decay_iters "$TRAIN_ITERS" \ --micro_batch_size "$MBS" \ --global_batch_size "$GBS" \ @@ -203,6 +314,7 @@ mkdir -p "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" --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" \ @@ -213,6 +325,8 @@ mkdir -p "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" --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 \ @@ -220,17 +334,18 @@ mkdir -p "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" --exp_avg_sq_dtype fp32 \ --enable_primus_turbo "$ENABLE_PRIMUS_TURBO" \ --use_turbo_attention "$USE_TURBO_ATTENTION" \ - --use_v4_triton_attention "$USE_V4_TRITON_ATTENTION" \ - --use_v4_triton_csa_attention "$USE_V4_TRITON_CSA_ATTENTION" \ - --use_v4_tilelang_attention "$USE_V4_TILELANG_ATTENTION" \ - --use_v4_tilelang_csa_attention "$USE_V4_TILELANG_CSA_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_mlp "$TURBO_USE_GROUPED_MLP" \ + --use_turbo_parallel_linear "$USE_TURBO_PARALLEL_LINEAR" \ + --moe_experts_fp4 "$MOE_EXPERTS_FP4" \ --moe_use_legacy_grouped_gemm False \ - --fp8 null \ - --fp8_recipe null \ + --fp8 "$FP8" \ + --fp8_recipe "$FP8_RECIPE" \ --recompute_num_layers 1 \ --recompute_granularity full \ --recompute_method uniform \ diff --git a/run_deepseek_v4_pro_muon_1gpu.sh b/run_deepseek_v4_pro_muon_1gpu.sh new file mode 100755 index 000000000..d00f071ef --- /dev/null +++ b/run_deepseek_v4_pro_muon_1gpu.sh @@ -0,0 +1,496 @@ +#!/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} +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/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_parallel_linear=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_mlp $TURBO_USE_GROUPED_MLP \ + --use_turbo_parallel_linear $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_TURBO_FUSE_GROUPED_WGRAD PRIMUS_TURBO_FUSE_WGRAD_DEBUG \ + 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/run_dsv4_projection_1gpu.sh b/run_dsv4_projection_1gpu.sh new file mode 100755 index 000000000..7deaa8c2e --- /dev/null +++ b/run_dsv4_projection_1gpu.sh @@ -0,0 +1,156 @@ +#!/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 + +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/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_mlp false --use_turbo_parallel_linear 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/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_muon_batched_orthogonalize.py b/tests/unit_tests/backends/megatron/test_muon_batched_orthogonalize.py new file mode 100644 index 000000000..bdb91875d --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_muon_batched_orthogonalize.py @@ -0,0 +1,52 @@ +# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. +"""Parity test for the batched grouped-expert Newton-Schulz fast path in +:meth:`TensorParallelMuon.orthogonalize`. + +At TP=1 the consolidated ``[E, N, K]`` expert momentum is orthogonalized in a +single batched Newton-Schulz call (``PRIMUS_MUON_BATCHED_NS`` default-on) +instead of an E-long per-expert Python loop. This asserts the batched path is +numerically equivalent to the loop it replaces (the loop is still used for the +partitioned TP>1 case, and is reachable via ``PRIMUS_MUON_BATCHED_NS=0``). +""" +from __future__ import annotations + +import pytest +import torch + +from primus.backends.megatron.core.optimizer.moun import ( + HAVE_EMERGING_OPTIMIZERS, + TensorParallelMuon, +) + +pytestmark = pytest.mark.skipif( + not HAVE_EMERGING_OPTIMIZERS, + reason="needs emerging_optimizers>=0.4.0a0 (batched 3-D newton_schulz)", +) + + +@pytest.mark.parametrize("shape", [(8, 64, 32), (6, 32, 48)]) +def test_batched_3d_orthogonalize_matches_loop(shape, monkeypatch): + """Batched 3-D orthogonalize == per-expert loop (CPU, fp32).""" + torch.manual_seed(0) + E, N, K = shape + + # fp32 + "high" matmul precision so the comparison isn't masked by the + # bf16 ("medium") Newton-Schulz path; both branches use the same precision, + # so the only difference under test is batched-vs-loop reduction order. + prev_prec = torch.get_float32_matmul_precision() + torch.set_float32_matmul_precision("high") + try: + p = torch.nn.Parameter(torch.zeros(E, N, K, dtype=torch.float32)) + opt = TensorParallelMuon([p], lr=1e-3, coefficient_type="deepseekv4", extra_scale_factor=0.18) + grad = torch.randn(E, N, K, dtype=torch.float32) + + monkeypatch.setenv("PRIMUS_MUON_BATCHED_NS", "1") + batched = opt.orthogonalize(p, grad.clone()) + + monkeypatch.setenv("PRIMUS_MUON_BATCHED_NS", "0") + loop = opt.orthogonalize(p, grad.clone()) + + assert batched.shape == (E, N, K) + torch.testing.assert_close(batched, loop, rtol=1e-4, atol=1e-4) + finally: + torch.set_float32_matmul_precision(prev_prec) 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..8aa510f3d --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_rope_arange_cache.py @@ -0,0 +1,47 @@ +# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. +"""Parity + caching test for ``RoPECache.forward_arange``. + +The compressed-branch RoPE is evaluated at the deterministic positions +``arange(P)`` every forward; ``forward_arange`` caches that table. This asserts +the cached table is bit-identical to recomputing ``forward(arange(n))`` and that +the cache actually memoises (and that ``PRIMUS_COMPRESS_ROPE_CACHE=0`` bypasses 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) diff --git a/tests/unit_tests/backends/megatron/test_v4_muon_param_groups.py b/tests/unit_tests/backends/megatron/test_v4_muon_param_groups.py index e70bc105a..c7aa5c997 100644 --- a/tests/unit_tests/backends/megatron/test_v4_muon_param_groups.py +++ b/tests/unit_tests/backends/megatron/test_v4_muon_param_groups.py @@ -17,7 +17,6 @@ import types -import pytest import torch from primus.backends.megatron.core.optimizer.moun import ( @@ -43,7 +42,7 @@ def test_2d_weight_matrices_go_to_muon(): # attention / MoE expert / mHC mixing matrix (HyperMixer.fn) are 2-D. assert _param_goes_to_muon(_p((4096, 1024))) is True # attn proj assert _param_goes_to_muon(_p((2048, 4096))) is True # MoE expert fc - assert _param_goes_to_muon(_p((24, 16384))) is True # HyperMixer.fn (K*D->(2+K)K) + assert _param_goes_to_muon(_p((24, 16384))) is True # HyperMixer.fn (K*D->(2+K)K) def test_embedding_and_output_go_to_adamw(): @@ -55,10 +54,10 @@ def test_embedding_and_output_go_to_adamw(): def test_1d_params_go_to_adamw(): # RMSNorm weight, mHC static bias (HyperMixer.base / HyperHead.base), # mHC gating scale (HyperMixer.scale shape [3]), and biases are 1-D. - assert _param_goes_to_muon(_p((4096,))) is False # RMSNorm weight - assert _param_goes_to_muon(_p((24,))) is False # HyperMixer.base - assert _param_goes_to_muon(_p((3,))) is False # HyperMixer.scale - assert _param_goes_to_muon(_p((4,))) is False # HyperHead.base (K) + assert _param_goes_to_muon(_p((4096,))) is False # RMSNorm weight + assert _param_goes_to_muon(_p((24,))) is False # HyperMixer.base + assert _param_goes_to_muon(_p((3,))) is False # HyperMixer.scale + assert _param_goes_to_muon(_p((4,))) is False # HyperHead.base (K) def test_scalar_params_go_to_adamw(): 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 index 55cc71790..afe006c24 100644 --- 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 @@ -744,16 +744,10 @@ def test_attention_spec_uses_column_and_row_parallel(): # 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.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.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. @@ -775,10 +769,7 @@ def test_attention_spec_uses_column_and_row_parallel(): # 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() - ) + assert submods_flat.linear_proj.module is provider.row_parallel_linear_with_scatter_input() def test_attention_spec_includes_compressor_and_indexer(): 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..41fdbe993 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_rope_from_positions.py @@ -0,0 +1,116 @@ +############################################################################### +# 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): + freqs = position_ids.float().unsqueeze(-1) * inv_freq + 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 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_p35_rope_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_p35_rope_triton.py index 63cb19e85..4a8461af5 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_p35_rope_triton.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_p35_rope_triton.py @@ -7,7 +7,7 @@ """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.rope_interleaved_partial``) +``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** @@ -48,7 +48,7 @@ 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.rope_interleaved_partial import ( # noqa: E402 +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, diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_p36_sinkhorn_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_p36_sinkhorn_triton.py index 09d461eaa..a37aa4d8c 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_p36_sinkhorn_triton.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_p36_sinkhorn_triton.py @@ -7,7 +7,7 @@ """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.sinkhorn``) +``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 @@ -57,7 +57,7 @@ _get_compiled_sinkhorn, sinkhorn_normalize, ) -from primus.backends.megatron.core.transformer.v4_attention_kernels._triton.sinkhorn import ( # noqa: E402 +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.sinkhorn import ( # noqa: E402 SinkhornNormalizeFn, eager_sinkhorn_normalize, is_triton_kernel_supported, diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_p37_hc_glue_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_p37_hc_glue_triton.py index 4f6ea7f0b..71c67c8c9 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_p37_hc_glue_triton.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_p37_hc_glue_triton.py @@ -7,7 +7,7 @@ """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.hc_glue``) +``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: @@ -43,7 +43,7 @@ from primus.backends.megatron.core.transformer.hyper_connection import ( # noqa: E402 HyperMixer, ) -from primus.backends.megatron.core.transformer.v4_attention_kernels._triton.hc_glue import ( # noqa: E402 +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, diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_p38_indexer_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_p38_indexer_triton.py index 84e55e94b..56379aeaf 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_p38_indexer_triton.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_p38_indexer_triton.py @@ -7,7 +7,7 @@ """Plan-6 P38 G41 — `Indexer.forward` scoring Triton parity. Asserts that :class:`IndexerScoreFn` (Triton kernel from -``primus.backends.megatron.core.transformer.v4_attention_kernels._triton.indexer_score``) +``primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score``) produces scores that match the eager body in :meth:`primus.backends.megatron.core.transformer.indexer.Indexer.forward` within bf16 tolerance, and that the load-bearing post-`topk` @@ -32,7 +32,7 @@ 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.indexer_score import ( # noqa: E402 +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score import ( # noqa: E402 IndexerScoreFn, is_triton_kernel_supported, ) @@ -47,7 +47,7 @@ # topk parity assertions still ran the right code. Switch the # import here to the post-tail module so the gate check matches # the env var the test sets. -from primus.backends.megatron.core.transformer.v4_attention_kernels._triton.indexer_score_post import ( # noqa: E402 +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score_post import ( # noqa: E402 is_triton_path_enabled, ) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_p39_router_post_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_p39_router_post_triton.py index 2f9210e47..d6f93a2f8 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_p39_router_post_triton.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_p39_router_post_triton.py @@ -33,7 +33,6 @@ from primus.backends.megatron.core.transformer.moe._triton.v4_router_post import ( # noqa: E402 V4RouterPostFn, - is_triton_kernel_supported, is_triton_path_enabled, ) from primus.backends.megatron.core.transformer.moe.v4_hash_router import ( # noqa: E402 @@ -123,14 +122,27 @@ def test_fast_tier_fwd_eager_parity(self, score_function, scale, out_dtype): 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 - K_pow2 = 8 # K must be power of 2 in the kernel; pad indices - logits, indices_full = _build_inputs(N=N, E=E, K=K_pow2, seed=8000) - # Use only the first K columns (K=6 unique indices per row). - indices = indices_full[:, :K_pow2] + 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 @@ -147,8 +159,9 @@ def test_release_tier_fwd_eager_parity(self, score_function): class TestG42BackwardParity: @pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) @pytest.mark.parametrize("scale", [1.0, 2.5]) - def test_fast_tier_bwd_eager_parity(self, score_function, scale): - N, E, K = 64, 32, 4 + @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) @@ -230,27 +243,25 @@ def test_hash_router_env_toggle_parity(self, score_function): class TestG42EdgeCases: - def test_unsupported_e_raises(self): - # E must be power of 2 - logits = torch.randn((4, 7), dtype=torch.float32, device="cuda") - indices = torch.zeros((4, 2), dtype=torch.int64, device="cuda") - with pytest.raises(ValueError, match="E must be"): - V4RouterPostFn.apply(logits, indices, "softmax", 1.0, torch.float32) - 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_supported_predicate(self): - good_logits = torch.randn((4, 8), dtype=torch.float32, device="cuda") - good_idx = torch.zeros((4, 2), dtype=torch.int64, device="cuda") - assert is_triton_kernel_supported(good_logits, good_idx) - - bad_e = torch.randn((4, 7), dtype=torch.float32, device="cuda") - assert not is_triton_kernel_supported(bad_e, good_idx) + 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) - assert not is_triton_kernel_supported(cpu_logits, cpu_idx) + 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_p41_indexer_tail_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_p41_indexer_tail_triton.py index ad327ed77..746583153 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_p41_indexer_tail_triton.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_p41_indexer_tail_triton.py @@ -7,7 +7,7 @@ """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.indexer_score_post``) +``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 @@ -35,7 +35,7 @@ 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.indexer_score_post import ( # 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, @@ -269,7 +269,7 @@ def test_env_explicit_zero_disables(self): 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.indexer_score import ( + from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score import ( is_triton_path_enabled as full_enabled, ) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_p49_tilelang_dispatch.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_p49_tilelang_dispatch.py index 6b6d002ee..c0766ad70 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_p49_tilelang_dispatch.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_p49_tilelang_dispatch.py @@ -249,7 +249,7 @@ def test_override_cache_dir(self): # --------------------------------------------------------------------------- -# G49.6: dispatcher wired into v4_attention / v4_csa_attention wrappers +# G49.6: dispatcher wired into v4_attention_v1 / v4_csa_attention_v0 wrappers # --------------------------------------------------------------------------- @@ -261,19 +261,19 @@ def test_v4_attention_calls_should_dispatch(self): import inspect mod = importlib.import_module( - "primus.backends.megatron.core.transformer.v4_attention_kernels.v4_attention" + "primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention" ) - src = inspect.getsource(mod.v4_attention) - assert "should_dispatch" in src, "v4_attention wrapper must call _tilelang.should_dispatch(...)" - assert "v4_attention_fwd_tilelang" in src, "v4_attention wrapper must route to the tilelang stub" + src = inspect.getsource(mod.v4_attention_v1) + assert "should_dispatch" in src, "v4_attention_v1 wrapper must call _tilelang.should_dispatch(...)" + assert "v4_attention_fwd_tilelang" in src, "v4_attention_v1 wrapper must route to the tilelang stub" def test_v4_csa_attention_calls_should_dispatch(self): import importlib import inspect mod = importlib.import_module( - "primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention" + "primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v0_deprecated.v4_csa_attention" ) - src = inspect.getsource(mod.v4_csa_attention) + src = inspect.getsource(mod.v4_csa_attention_v0) assert "should_dispatch" in src assert "v4_csa_attention_fwd_tilelang" in src diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_p50_v4_attention_fwd_tilelang.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_p50_v4_attention_fwd_tilelang.py index d46e0b85a..b35d006ad 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_p50_v4_attention_fwd_tilelang.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_p50_v4_attention_fwd_tilelang.py @@ -31,12 +31,12 @@ # environment instead of erroring at import time. pytest.importorskip("tilelang.language", reason="tilelang not installed") +from primus.backends.megatron.core.transformer.v4_attention_kernels._eager.reference import ( # noqa: E402 + eager_v4_attention, +) from primus.backends.megatron.core.transformer.v4_attention_kernels._tilelang.v4_attention_fwd_tilelang import ( # noqa: E402 v4_attention_fwd_tilelang, ) -from primus.backends.megatron.core.transformer.v4_attention_kernels.reference import ( # noqa: E402 - eager_v4_attention, -) def _build_inputs( diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_p51_v4_attention_bwd_tilelang.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_p51_v4_attention_bwd_tilelang.py index 05b0c7e3f..d50e0cbd7 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_p51_v4_attention_bwd_tilelang.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_p51_v4_attention_bwd_tilelang.py @@ -30,13 +30,13 @@ # from the vendored source tree. pytest.importorskip("tilelang.language", reason="tilelang not installed") +from primus.backends.megatron.core.transformer.v4_attention_kernels._eager.reference import ( # noqa: E402 + eager_v4_attention, +) from primus.backends.megatron.core.transformer.v4_attention_kernels._tilelang.v4_attention_autograd_tilelang import ( # noqa: E402 V4AttentionTilelangFn, v4_attention_tilelang, ) -from primus.backends.megatron.core.transformer.v4_attention_kernels.reference import ( # noqa: E402 - eager_v4_attention, -) def _build_inputs(*, B, HQ, HK, Sq, Sk, D, dtype, seed=20260515, has_sink=True): @@ -157,20 +157,20 @@ def test_bwd_kernel_registered(self): def test_full_dispatch_uses_autograd_path(self): """When ``use_tilelang=True`` and both FWD/BWD are registered, the - public ``v4_attention()`` should route through ``V4AttentionTilelangFn``. + public ``v4_attention_v1()`` should route through ``V4AttentionTilelangFn``. Plan-8 P57 close-out 2: the legacy ``PRIMUS_V4_TILELANG_ATTN`` env knob is gone; the caller now plumbs the config flag via the ``use_tilelang`` kwarg. """ from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_attention import ( - v4_attention, + v4_attention_v1, ) B, HQ, Sq, Sk, D = 1, 4, 32, 32, 64 q, k, v, sink = _build_inputs(B=B, HQ=HQ, HK=1, Sq=Sq, Sk=Sk, D=D, dtype=torch.bfloat16) q.requires_grad_(True) - out = v4_attention( + out = v4_attention_v1( q, k, v, 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..fc2d97a5c --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_backend_import_gating.py @@ -0,0 +1,115 @@ +############################################################################### +# 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") + 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_flydsl_csa_attention.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_flydsl_csa_attention.py new file mode 100644 index 000000000..3194f1c98 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_flydsl_csa_attention.py @@ -0,0 +1,366 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""FlyDSL CSA (compress_ratio == 4) attention fwd / bwd correctness. + +This module checks the **FlyDSL** CSA kernels against the eager-Python +reference (:func:`eager_v4_csa_attention`, the bit-exact "truth" defined +in ``v4_attention_kernels/reference.py``) and, for context, reports the +in-tree **Triton** kernel error against the same golden so the two +backends can be compared side by side. + +Scope / constraints (driven by the FlyDSL CSA kernel preconditions): + +* ``compress_ratio == 4`` (CSA) only. +* bf16 only, ``head_dim == 512`` (``D % 64 == 0`` for bwd), ``swa_window + > 0``, ``scale == 1/sqrt(D)`` — these mirror the asserts inside + ``_launch_v4_attention_fwd_csa`` / ``flydsl_v4_csa_attention_bwd``. +* K/V are MQA single-latent (head-identical): the FlyDSL CSA backward + reads head 0 only (``k_local[:, :1]``), so the inputs are built by + broadcasting a ``[B, 1, S, D]`` latent across heads — exactly what the + V4 forward feeds the kernel in production. +* The FlyDSL backward only runs its own kernel when + ``V4_FLYDSL_CSA_BWD_FLY_DQ=1`` and ``V4_FLYDSL_CSA_BWD_FLY_DKV=1`` + (otherwise the launcher transparently falls back to Triton). The bwd + test sets both knobs so the FlyDSL kernel is actually exercised. + +The golden reference is evaluated in fp32 on the **same bf16-rounded +input bits** (``x.float()``) so the measured error reflects kernel +compute precision (bf16 tensor-core matmul + fp32 softmax) rather than +input-quantisation noise. + +GPU-only; skipped at collection time on CPU or when FlyDSL is absent. +""" + +from __future__ import annotations + +import math +from typing import Tuple + +import pytest + +torch = pytest.importorskip("torch") + +if not torch.cuda.is_available(): + pytest.skip("FlyDSL CSA kernels require CUDA / HIP", allow_module_level=True) + +pytest.importorskip("triton", reason="Triton baseline not installed") + +# Import the local primus reference + Triton launchers FIRST so they are +# cached in sys.modules before the FlyDSL kernel wrappers (which insert +# ``/workspace/Primus`` onto sys.path at import time) can shadow them. +from primus.backends.megatron.core.transformer.v4_attention_kernels import ( # noqa: E402 + eager_v4_csa_attention, +) +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_v0_deprecated.v4_csa_attention_fwd import ( # noqa: E402 + _launch_v4_csa_attention_fwd, +) + +try: + from primus.backends.megatron.core.transformer.v4_attention_kernels._flydsl_v0_deprecated.kernels.v4_attention_fwd_flydsl_csa import ( # noqa: E402 + _launch_v4_attention_fwd_csa, + ) + from primus.backends.megatron.core.transformer.v4_attention_kernels._flydsl_v0_deprecated.kernels.v4_csa_attention_bwd_flydsl_mqa import ( # noqa: E402 + flydsl_v4_csa_attention_bwd, + ) +except Exception as exc: # pragma: no cover - environment-dependent + pytest.skip(f"FlyDSL CSA kernels unavailable: {exc!r}", allow_module_level=True) + + +# --------------------------------------------------------------------------- +# Shapes — small (fp32 golden must fit HBM) but with the production +# head_dim=512 that FlyDSL hard-requires. +# --------------------------------------------------------------------------- + +# (variant, B, H, S, D, K_topk, swa_window) +_SHAPES = [ + ("flash_like", 1, 8, 256, 512, 128, 128), + ("pro_like", 1, 16, 256, 512, 128, 128), +] +_SINK_MODES = [True, False] + + +def _make_inputs( + *, + B: int, + H: int, + S: int, + D: int, + K_topk: int, + sink_on: bool, + seed: int = 1234, +): + """Build bf16 CSA inputs with MQA (head-identical) K/V. + + Returns the bf16 tensors the kernels consume plus the scalar config. + ``sparse_mask`` drops ~25% of slots (``-inf``) and the matching + ``gathered`` rows are zeroed so eager / Triton / FlyDSL all see the + same physical data on the masked positions. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + device = "cuda" + dtype = torch.bfloat16 + + q = torch.randn(B, H, S, D, generator=g, device=device, dtype=dtype) + # MQA single latent. ``k_mqa`` / ``v_mqa`` are the [B, 1, S, D] tensors + # the FlyDSL forward consumes (it uses the mqa_kv stride trick); the + # [B, H, S, D] broadcast (``k_full`` / ``v_full``) is what eager / Triton + # consume and is the autograd leaf, so per-head grads can be compared + # against the FlyDSL [B, H, S, D] dk/dv buffers. Heads are identical, so + # all three backends see the same physical K/V. + k_mqa = torch.randn(B, 1, S, D, generator=g, device=device, dtype=dtype) + v_mqa = torch.randn(B, 1, S, D, generator=g, device=device, dtype=dtype) + k_full = k_mqa.expand(B, H, S, D).contiguous() + v_full = v_mqa.expand(B, H, S, D).contiguous() + + gathered = torch.randn(B, S, K_topk, D, generator=g, device=device, dtype=dtype) + valid = torch.rand(B, S, K_topk, generator=g, device=device) > 0.25 + sparse_mask = torch.where( + valid, + torch.zeros((), dtype=dtype, device=device), + torch.tensor(float("-inf"), dtype=dtype, device=device), + ) + gathered = gathered * valid.unsqueeze(-1).to(dtype) + + sink = torch.randn(H, generator=g, device=device, dtype=torch.float32) * 0.1 if sink_on else None + return q, k_full, v_full, k_mqa, v_mqa, gathered, sparse_mask, sink + + +def _err(cand: torch.Tensor, ref: torch.Tensor) -> Tuple[float, float]: + """Max abs error and max relative error (vs |ref|, eps-floored).""" + cand = cand.float() + ref = ref.float() + abs_err = (cand - ref).abs() + max_abs = abs_err.max().item() + denom = ref.abs().clamp_min(1e-6) + max_rel = (abs_err / denom).max().item() + return max_abs, max_rel + + +# --------------------------------------------------------------------------- +# Forward correctness +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("variant,B,H,S,D,K_topk,swa", _SHAPES, ids=[s[0] for s in _SHAPES]) +@pytest.mark.parametrize("sink_on", _SINK_MODES, ids=["sink_on", "sink_off"]) +def test_flydsl_csa_fwd_matches_eager( + variant: str, + B: int, + H: int, + S: int, + D: int, + K_topk: int, + swa: int, + sink_on: bool, +): + """FlyDSL CSA forward vs fp32 eager golden (Triton reported for context).""" + q, k_full, v_full, k_mqa, v_mqa, gathered, sparse_mask, sink = _make_inputs( + B=B, H=H, S=S, D=D, K_topk=K_topk, sink_on=sink_on + ) + scale = 1.0 / math.sqrt(D) + + # fp32 golden on the same bf16-rounded bits. + out_gold = eager_v4_csa_attention( + q.float(), + k_full.float(), + v_full.float(), + gathered.float(), + sink=None if sink is None else sink.float(), + swa_window=swa, + sparse_mask=sparse_mask.float(), + attn_dropout=0.0, + training=False, + scale=scale, + ) + + # FlyDSL forward consumes the MQA [B, 1, S, D] latent (mqa_kv path). + out_fly, _lse_fly = _launch_v4_attention_fwd_csa( + q, + k_mqa, + v_mqa, + gathered, + sink=sink, + swa_window=swa, + sparse_mask=sparse_mask, + scale=scale, + ) + # Triton consumes the [B, H, S, D] broadcast (its production input). + out_tri, _lse_tri = _launch_v4_csa_attention_fwd( + q, + k_full, + v_full, + gathered, + sparse_mask, + sink=sink, + swa_window=swa, + scale=scale, + ) + + fly_abs, fly_rel = _err(out_fly, out_gold) + tri_abs, tri_rel = _err(out_tri, out_gold) + print( + f"\n[CSA fwd] {variant} sink={sink_on} H={H} S={S} D={D} K={K_topk}\n" + f" FlyDSL vs golden: max_abs={fly_abs:.3e} max_rel={fly_rel:.3e}\n" + f" Triton vs golden: max_abs={tri_abs:.3e} max_rel={tri_rel:.3e}", + flush=True, + ) + + assert out_fly.shape == out_gold.shape == q.shape + assert out_fly.dtype == torch.bfloat16 + # bf16 tolerance; Triton must pass (sanity), FlyDSL failing is the signal. + torch.testing.assert_close(out_tri.float(), out_gold, atol=2e-2, rtol=2e-2) + torch.testing.assert_close(out_fly.float(), out_gold, atol=2e-2, rtol=2e-2) + + +# --------------------------------------------------------------------------- +# Backward correctness +# --------------------------------------------------------------------------- + + +def _eager_grads(q, k_local, v_local, gathered, sparse_mask, sink, *, swa, scale, dout): + """fp32 eager autograd grads on the same bf16-rounded input bits.""" + qg = q.float().detach().requires_grad_(True) + kg = k_local.float().detach().requires_grad_(True) + vg = v_local.float().detach().requires_grad_(True) + gg = gathered.float().detach().requires_grad_(True) + sg = None if sink is None else sink.float().detach().requires_grad_(True) + + out = eager_v4_csa_attention( + qg, + kg, + vg, + gg, + sink=sg, + swa_window=swa, + sparse_mask=sparse_mask.float(), + attn_dropout=0.0, + training=False, + scale=scale, + ) + out.backward(dout.float()) + return qg.grad, kg.grad, vg.grad, gg.grad, (None if sg is None else sg.grad) + + +@pytest.mark.parametrize("variant,B,H,S,D,K_topk,swa", _SHAPES, ids=[s[0] for s in _SHAPES]) +@pytest.mark.parametrize("sink_on", _SINK_MODES, ids=["sink_on", "sink_off"]) +def test_flydsl_csa_bwd_matches_eager( + variant: str, + B: int, + H: int, + S: int, + D: int, + K_topk: int, + swa: int, + sink_on: bool, + monkeypatch, +): + """FlyDSL CSA backward vs fp32 eager golden (Triton reported for context). + + Forces ``V4_FLYDSL_CSA_BWD_FLY_DQ=1`` + ``V4_FLYDSL_CSA_BWD_FLY_DKV=1`` + so a single FlyDSL full-kernel launch produces all five grads (else + the launcher would fall back to Triton). + """ + q, k_full, v_full, k_mqa, v_mqa, gathered, sparse_mask, sink = _make_inputs( + B=B, H=H, S=S, D=D, K_topk=K_topk, sink_on=sink_on, seed=2025 + ) + scale = 1.0 / math.sqrt(D) + gen = torch.Generator(device="cuda").manual_seed(99) + dout = torch.randn(B, H, S, D, generator=gen, device="cuda", dtype=torch.bfloat16) + + # Golden grads (eager autograd over the [B, H, S, D] leaves). + dq_g, dk_g, dv_g, dg_g, ds_g = _eager_grads( + q, k_full, v_full, gathered, sparse_mask, sink, swa=swa, scale=scale, dout=dout + ) + + # FlyDSL fwd (MQA latent) -> out + lse feed FlyDSL bwd. The bwd consumes + # the [B, H, S, D] tensor and slices head 0 internally (mqa_kv=True). + out_fly, lse_fly = _launch_v4_attention_fwd_csa( + q, + k_mqa, + v_mqa, + gathered, + sink=sink, + swa_window=swa, + sparse_mask=sparse_mask, + scale=scale, + ) + monkeypatch.setenv("V4_FLYDSL_CSA_BWD_FLY_DQ", "1") + monkeypatch.setenv("V4_FLYDSL_CSA_BWD_FLY_DKV", "1") + dq_f, dk_f, dv_f, dg_f, ds_f = flydsl_v4_csa_attention_bwd( + q, + k_full, + v_full, + gathered, + sparse_mask, + out_fly, + dout, + lse_fly, + sink=sink, + swa_window=swa, + scale=scale, + ) + + # Triton baseline (its own fwd out+lse -> its own bwd). + out_tri, lse_tri = _launch_v4_csa_attention_fwd( + q, + k_full, + v_full, + gathered, + sparse_mask, + sink=sink, + swa_window=swa, + scale=scale, + ) + dq_t, dk_t, dv_t, dg_t, ds_t = _launch_v4_csa_attention_bwd( + q, + k_full, + v_full, + gathered, + sparse_mask, + out_tri, + dout, + lse_tri, + sink=sink, + swa_window=swa, + scale=scale, + ) + + def _report(name, fly, tri, gold): + fa, fr = _err(fly, gold) + ta, tr = _err(tri, gold) + print( + f" {name:10s} FlyDSL max_abs={fa:.3e} max_rel={fr:.3e} | " + f"Triton max_abs={ta:.3e} max_rel={tr:.3e}", + flush=True, + ) + return fa, fr + + print( + f"\n[CSA bwd] {variant} sink={sink_on} H={H} S={S} D={D} K={K_topk}", + flush=True, + ) + _report("dq", dq_f, dq_t, dq_g) + _report("dk_local", dk_f, dk_t, dk_g) + _report("dv_local", dv_f, dv_t, dv_g) + _report("dgathered", dg_f, dg_t, dg_g) + if sink_on: + _report("dsink", ds_f, ds_t, ds_g) + + # bf16 gradient tolerance (looser than fwd: grads accumulate more + # terms). Triton must pass as a sanity floor; FlyDSL failing flags a + # precision problem to report (do NOT auto-fix). + bwd_tol = dict(atol=5e-2, rtol=5e-2) + torch.testing.assert_close(dq_t.float(), dq_g, **bwd_tol) + torch.testing.assert_close(dq_f.float(), dq_g, **bwd_tol) + torch.testing.assert_close(dk_f.float(), dk_g, **bwd_tol) + torch.testing.assert_close(dv_f.float(), dv_g, **bwd_tol) + torch.testing.assert_close(dg_f.float(), dg_g, **bwd_tol) + if sink_on: + torch.testing.assert_close(ds_f.float(), ds_g, **bwd_tol) 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_mtp.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_mtp.py index 8445732b4..fff550143 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_mtp.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_mtp.py @@ -58,9 +58,6 @@ from megatron.core.transformer.spec_utils import ModuleSpec from megatron.core.transformer.transformer_layer import TransformerLayerSubmodules -from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import ( - DeepSeekV4SpecProvider, -) from primus.backends.megatron.core.models.deepseek_v4 import ( DeepseekV4HybridLayer, DeepseekV4HybridLayerSubmodules, diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p21_strict_build.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p21_strict_build.py index 5c12786c5..52f03b60d 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p21_strict_build.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p21_strict_build.py @@ -124,8 +124,7 @@ def test_no_try_except_returns_nn_linear(self): "provider-built module instantiates cleanly (see " "DeepSeekV4SpecProvider.column_parallel_linear_with_gather_output " "/ row_parallel_linear_with_scatter_input for the canonical " - "non-TE path), or let build_module raise.\n\nOffenders:\n " - + "\n ".join(offenders) + "non-TE path), or let build_module raise.\n\nOffenders:\n " + "\n ".join(offenders) ) @pytest.mark.parametrize( @@ -273,7 +272,6 @@ def _tp1_distributed(self): """Initialize a 1-rank torch.distributed group on gloo (CPU).""" import os - import torch import torch.distributed as dist from megatron.core import parallel_state @@ -301,9 +299,6 @@ def _tp1_distributed(self): def test_v4_attention_strict_build_at_tp1(self, _tp1_distributed): import torch - from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( - DeepSeekV4TransformerConfig, - ) # Order matters: importing the V4 model package above (via the # config import) makes ``deepseek_v4_block`` finish loading, # which then unblocks the ``deepseek_v4_attention`` import below. @@ -313,6 +308,9 @@ def test_v4_attention_strict_build_at_tp1(self, _tp1_distributed): from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_layer_specs import ( _build_v4_attention_submodules, ) + 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, ) @@ -362,13 +360,15 @@ def test_v4_attention_strict_build_at_tp1(self, _tp1_distributed): ) # Sanity: every linear is a Megatron parallel module, not a bare nn.Linear. + from megatron.core.extensions.transformer_engine import TELinear from megatron.core.tensor_parallel.layers import ( ColumnParallelLinear, RowParallelLinear, ) - from primus.backends.megatron.core.extensions.primus_turbo import PrimusTurboLinear - from megatron.core.extensions.transformer_engine import TELinear + from primus.backends.megatron.core.extensions.primus_turbo import ( + PrimusTurboLinear, + ) sharded_or_te = ( ColumnParallelLinear, diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p22_core_attention.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p22_core_attention.py index 1eb5690fc..2f4505008 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p22_core_attention.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p22_core_attention.py @@ -39,10 +39,8 @@ class that replicates the eager-Python scaled-dot-product math is import pytest import torch - from megatron.core.transformer.enums import AttnMaskType -from megatron.core.transformer.spec_utils import ModuleSpec, build_module - +from megatron.core.transformer.spec_utils import ModuleSpec # --------------------------------------------------------------------------- # G18 — submodules surface @@ -153,9 +151,7 @@ class TestSpecEmission: 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." - ) + 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) @@ -321,9 +317,7 @@ def test_no_alias_when_core_attention_lacks_sink(self, _tp1_distributed): ) provider = DeepSeekV4SpecProvider(config=cfg) - attn2 = _build_v4_attention_with( - cfg, core_attention_class=provider.core_attention() - ) + 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 " @@ -341,9 +335,7 @@ def test_dense_forward_matches_eager(self, _tp1_distributed): 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 - ) + attn_turbo = _build_v4_attention_with(cfg, core_attention_class=_MockTurboCoreAttention) # Copy weights so both modules have identical parameters. with torch.no_grad(): diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p23_turbo_deepep_dispatcher.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p23_turbo_deepep_dispatcher.py index 62a2200b6..2ced05fac 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p23_turbo_deepep_dispatcher.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p23_turbo_deepep_dispatcher.py @@ -58,7 +58,6 @@ from unittest.mock import MagicMock import pytest - from megatron.core.transformer.moe.token_dispatcher import ( MoEAllGatherTokenDispatcher, MoEAlltoAllTokenDispatcher, @@ -76,7 +75,6 @@ ) from primus.backends.megatron.core.transformer.moe.v4_moe import DeepseekV4MoE - # --------------------------------------------------------------------------- # Test fixtures # --------------------------------------------------------------------------- @@ -115,9 +113,7 @@ def _fake(name): return MagicMock() return real_find_spec(name) - monkeypatch.setattr( - deepseek_v4_layer_specs.importlib.util, "find_spec", _fake - ) + monkeypatch.setattr(deepseek_v4_layer_specs.importlib.util, "find_spec", _fake) @pytest.fixture @@ -130,9 +126,7 @@ def _fake(name): return None return real_find_spec(name) - monkeypatch.setattr( - deepseek_v4_layer_specs.importlib.util, "find_spec", _fake - ) + monkeypatch.setattr(deepseek_v4_layer_specs.importlib.util, "find_spec", _fake) # ``type(..., (), {})`` produces a class whose ``__name__`` is the @@ -293,9 +287,7 @@ def test_flex_without_turbo(self): assert cls is MoEFlexTokenDispatcher assert type_name == "flex" - def test_flex_with_turbo_active( - self, primus_turbo_available, fake_turbo_class - ): + 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, @@ -306,9 +298,7 @@ def test_flex_with_turbo_active( assert cls is fake_turbo_class assert type_name == "flex" - def test_flex_with_turbo_active_but_class_missing( - self, primus_turbo_available, monkeypatch - ): + 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, @@ -325,9 +315,7 @@ def test_flex_with_turbo_active_but_class_missing( assert cls is MoEFlexTokenDispatcher assert type_name == "flex" - def test_flex_with_turbo_inactive_tp_gt_1( - self, primus_turbo_available, fake_turbo_class - ): + 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( @@ -345,10 +333,7 @@ def test_unknown_type_falls_back_to_alltoall(self, caplog): 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 - ) + 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 @@ -379,23 +364,16 @@ class TestV4MoEResolver: def test_turbo_class_resolves_to_flex(self): spec = ModuleSpec(module=_FakeTurboDispatcher) - assert ( - DeepseekV4MoE._resolve_dispatcher_type_from_spec(spec) == "flex" - ) + 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" - ) + 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" - ) + assert DeepseekV4MoE._resolve_dispatcher_type_from_spec(spec) == "alltoall" def test_flex_unchanged(self): spec = ModuleSpec(module=MoEFlexTokenDispatcher) @@ -403,9 +381,7 @@ def test_flex_unchanged(self): def test_allgather_unchanged(self): spec = ModuleSpec(module=MoEAllGatherTokenDispatcher) - assert ( - DeepseekV4MoE._resolve_dispatcher_type_from_spec(spec) == "allgather" - ) + assert DeepseekV4MoE._resolve_dispatcher_type_from_spec(spec) == "allgather" def test_unknown_class_falls_back_to_alltoall(self, caplog): class _Unknown: @@ -415,6 +391,4 @@ class _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 - ) + 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_p25_dispatch.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p25_dispatch.py index 08d82d15d..f9c6da1bc 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p25_dispatch.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p25_dispatch.py @@ -4,29 +4,21 @@ # See LICENSE for license information. ############################################################################### -"""Plan-4 P25 \u2014 dispatch precedence in :class:`DeepseekV4Attention`. - -Asserts the static plumbing that wires ``use_v4_triton_attention`` -into the V4 attention forward path, without depending on a full -forward pass: - -* :class:`DeepSeekV4TransformerConfig` exposes - ``use_v4_triton_attention: bool`` and - ``use_v4_triton_csa_attention: bool``, -* :class:`DeepseekV4Attention.__init__` reads - ``config.use_v4_triton_attention`` and stores it as - ``self._use_v4_triton_attention``, -* the flag is auto-disabled when ``compress_ratio == 4`` (CSA path - uses the separate ``use_v4_triton_csa_attention`` flag landing in - P26), -* the dense / HCA forward dispatch in ``forward`` reads - ``self._use_v4_triton_attention`` after the - ``self._use_core_attention`` Turbo branch, enforcing the documented - precedence ``use_turbo_attention > use_v4_triton_attention > eager``. - -The full kernel-vs-eager output equivalence is covered by G23 / G24 -(``test_v4_p25_v4_attention_fwd.py`` / ``test_v4_p25_v4_attention_bwd.py``); -this file is the static / wiring-side gate. +"""Dense/HCA attention dispatch plumbing in :class:`DeepseekV4Attention`. + +Static (wiring) gate for the unified ``use_v4_attention_backend`` selector: + +* :class:`DeepSeekV4TransformerConfig` exposes ``use_v4_attention_backend: str`` + (and ``use_v4_csa_attention_backend``) defaulting to ``"triton_v1"``; +* :class:`DeepseekV4Attention.__init__` reads them into ``self._attn_backend`` / + ``self._csa_backend``; +* the dense/HCA dispatch goes through ``_attention_backend_forward``, which + switches on ``self._attn_backend`` (``use_turbo_attention`` still wins via the + earlier ``_use_core_attention`` branch), and the ``triton_v1`` path calls + ``v4_attention_v1`` through ``_attention_forward_via_v4_triton``. + +Output equivalence is covered by the G23/G24 fwd/bwd tests; this is the static +wiring gate. """ from __future__ import annotations @@ -35,96 +27,61 @@ DeepSeekV4TransformerConfig, ) -# --------------------------------------------------------------------------- -# Static plumbing — config + module surface -# --------------------------------------------------------------------------- - - -def test_p25_config_exposes_use_v4_triton_attention(): - """V4 transformer config exposes both Triton-attention switches.""" - fields = {f.name for f in DeepSeekV4TransformerConfig.__dataclass_fields__.values()} - assert "use_v4_triton_attention" in fields - assert "use_v4_triton_csa_attention" in fields - # Defaults must be False so existing checkpoints / smokes are - # unaffected by the new switches. - cfg_field = DeepSeekV4TransformerConfig.__dataclass_fields__["use_v4_triton_attention"] - csa_field = DeepSeekV4TransformerConfig.__dataclass_fields__["use_v4_triton_csa_attention"] - assert cfg_field.default is False - assert csa_field.default is False +def test_config_exposes_backend_selectors(): + """V4 transformer config exposes the unified string selectors, default triton_v1.""" + fields = DeepSeekV4TransformerConfig.__dataclass_fields__ + assert "use_v4_attention_backend" in fields + assert "use_v4_csa_attention_backend" in fields + assert fields["use_v4_attention_backend"].default == "triton_v1" + assert fields["use_v4_csa_attention_backend"].default == "triton_v1" -def test_p25_attention_init_reads_use_v4_triton_attention_flag(): - """``DeepseekV4Attention.__init__`` reads the config flag once.""" +def test_attention_init_reads_backend_selector(): + """``DeepseekV4Attention.__init__`` reads config.use_v4_attention_backend.""" from primus.backends.megatron.core.transformer.deepseek_v4_attention import ( DeepseekV4Attention, ) - src = DeepseekV4Attention.__init__.__code__.co_consts - # The init body references the flag name as a string constant when - # ``getattr(config, "use_v4_triton_attention", False)`` is compiled. - assert "use_v4_triton_attention" in src, ( - "DeepseekV4Attention.__init__ must read config.use_v4_triton_attention " "(plan-4 P25 plumbing)." - ) + init_consts = " ".join(str(c) for c in DeepseekV4Attention.__init__.__code__.co_consts if c is not None) + assert "use_v4_attention_backend" in init_consts + assert "use_v4_csa_attention_backend" in init_consts -def test_p25_forward_consults_use_v4_triton_attention_for_dense_and_hca(): - """``forward``'s cr == 0 / cr == 128 branches gate on the flag.""" +def test_backend_forward_helper_switches_on_attn_backend(): + """The dense/HCA dispatch helper switches on ``self._attn_backend``.""" from primus.backends.megatron.core.transformer.deepseek_v4_attention import ( DeepseekV4Attention, ) - forward_src = DeepseekV4Attention.forward.__code__.co_consts - # Must reference the runtime flag attribute name. ``co_consts`` may - # not surface attribute strings directly when Python compiles a - # ``self._use_v4_triton_attention`` access (the attribute name lives - # in ``co_names``); check both. - names = set(DeepseekV4Attention.forward.__code__.co_names) - has_flag_ref = "_use_v4_triton_attention" in names or any( - "_use_v4_triton_attention" in str(c) for c in forward_src - ) - assert has_flag_ref, ( - "DeepseekV4Attention.forward must consult self._use_v4_triton_attention " - "(plan-4 P25 dispatch — precedence " - "use_turbo_attention > use_v4_triton_attention > eager)." - ) + assert hasattr(DeepseekV4Attention, "_attention_backend_forward") + co = DeepseekV4Attention._attention_backend_forward.__code__ + names = set(co.co_names) + consts = " ".join(str(c) for c in co.co_consts if c is not None) + assert "_attn_backend" in names + # references each supported dense/HCA backend value + its entry + for be in ("gluon", "triton_v2", "triton_v1"): + assert be in consts, f"helper must handle backend value {be!r}" -def test_p25_forward_helper_method_exists(): - """Plan-4 P25 ships a dedicated forward helper for the Triton path.""" +def test_forward_uses_backend_forward_helper(): + """``forward`` routes the dense/HCA paths through the helper.""" from primus.backends.megatron.core.transformer.deepseek_v4_attention import ( DeepseekV4Attention, ) - assert hasattr(DeepseekV4Attention, "_attention_forward_via_v4_triton") - helper = DeepseekV4Attention._attention_forward_via_v4_triton - # The helper must call the Triton entry point (string ref in code). - co = helper.__code__ - names = set(co.co_names) - consts_str = " ".join(str(c) for c in co.co_consts if c is not None) - refs_v4_attention = "v4_attention" in names or "v4_attention" in consts_str - assert refs_v4_attention, "_attention_forward_via_v4_triton must call v4_attention (plan-4 P25)." - - -def test_p25_csa_layers_disable_v4_triton_attention_flag(): - """CSA (cr == 4) layers MUST NOT pick up the dense/HCA Triton flag. + names = set(DeepseekV4Attention.forward.__code__.co_names) + assert "_attention_backend_forward" in names - The dispatch is: - cr ∈ {0, 128} ← ``use_v4_triton_attention`` - cr == 4 ← ``use_v4_triton_csa_attention`` (plan-4 P26) - The init's auto-disable branch surfaces this contract explicitly so - a stray run script with ``use_v4_triton_attention=True`` does not - silently accelerate cr ∈ {0, 128} while leaving CSA on eager and - skewing apples-to-apples perf comparisons. - """ +def test_triton_v1_helper_calls_v4_attention_v1(): + """The triton_v1 path still calls ``v4_attention_v1`` via the launcher helper.""" from primus.backends.megatron.core.transformer.deepseek_v4_attention import ( DeepseekV4Attention, ) - # The init body contains an early-return / disable branch that - # references both the flag name and the cr-not-in-(0, 128) check. - init_consts = " ".join(str(c) for c in DeepseekV4Attention.__init__.__code__.co_consts if c is not None) - assert "use_v4_triton_attention" in init_consts - # The (0, 128) tuple is the constexpr that gates the auto-disable. - assert (0, 128) in DeepseekV4Attention.__init__.__code__.co_consts + assert hasattr(DeepseekV4Attention, "_attention_forward_via_v4_triton") + co = DeepseekV4Attention._attention_forward_via_v4_triton.__code__ + names = set(co.co_names) + consts = " ".join(str(c) for c in co.co_consts if c is not None) + assert "v4_attention_v1" in names or "v4_attention_v1" in consts diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p25_v4_attention_bwd.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p25_v4_attention_bwd.py index 4803d2f5a..302ed8f38 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p25_v4_attention_bwd.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p25_v4_attention_bwd.py @@ -4,10 +4,10 @@ # See LICENSE for license information. ############################################################################### -"""Plan-4 P25 G24 — `v4_attention` Triton BWD equivalence to autograd-on-eager. +"""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`'s autograd Function match the gradients computed +: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 / @@ -27,7 +27,7 @@ torch = pytest.importorskip("torch") if not torch.cuda.is_available(): - pytest.skip("v4_attention Triton kernel requires CUDA / HIP", allow_module_level=True) + pytest.skip("v4_attention_v1 Triton kernel requires CUDA / HIP", allow_module_level=True) pytest.importorskip("triton", reason="Triton not installed") @@ -36,7 +36,7 @@ ) from primus.backends.megatron.core.transformer.v4_attention_kernels import ( # noqa: E402 eager_v4_attention, - v4_attention, + v4_attention_v1, ) # --------------------------------------------------------------------------- @@ -281,7 +281,7 @@ def test_g24_dense_bwd_matches_eager( out_ref, ref_inp["q"], ref_inp["k"], ref_inp["v"], ref_inp["sink"] ) - out_cand = v4_attention( + out_cand = v4_attention_v1( cand_inp["q"], cand_inp["k"], cand_inp["v"], @@ -373,7 +373,7 @@ def test_g24_hca_style_bwd_matches_eager( out_ref, ref_inp["q"], ref_inp["k"], ref_inp["v"], ref_inp["sink"] ) - out_cand = v4_attention( + out_cand = v4_attention_v1( cand_inp["q"], cand_inp["k"], cand_inp["v"], @@ -420,7 +420,7 @@ def test_g24_no_nan_in_grads_dense_fp32(): use_hca=False, seed=4321, ) - out_cand = v4_attention( + out_cand = v4_attention_v1( inp["q"], inp["k"], inp["v"], diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p25_v4_attention_fwd.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p25_v4_attention_fwd.py index 80173be9b..d28d0cc59 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p25_v4_attention_fwd.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p25_v4_attention_fwd.py @@ -4,9 +4,9 @@ # See LICENSE for license information. ############################################################################### -"""Plan-4 P25 G23 — `v4_attention` Triton FWD equivalence to eager. +"""Plan-4 P25 G23 — `v4_attention_v1` Triton FWD equivalence to eager. -Asserts that :func:`v4_attention` (Triton kernel from +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: @@ -31,7 +31,7 @@ torch = pytest.importorskip("torch") if not torch.cuda.is_available(): - pytest.skip("v4_attention Triton kernel requires CUDA / HIP", allow_module_level=True) + 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") @@ -41,7 +41,7 @@ ) from primus.backends.megatron.core.transformer.v4_attention_kernels import ( # noqa: E402 eager_v4_attention, - v4_attention, + v4_attention_v1, ) # --------------------------------------------------------------------------- @@ -246,8 +246,8 @@ def test_g23_dense_fwd_matches_eager( scale=scale, ) - # Candidate: v4_attention with swa_window > 0, additive_mask=None - out_cand = v4_attention( + # Candidate: v4_attention_v1 with swa_window > 0, additive_mask=None + out_cand = v4_attention_v1( toy["q"], toy["k"], toy["v"], @@ -316,7 +316,7 @@ def test_g23_hca_style_fwd_matches_eager( scale=scale, ) - out_cand = v4_attention( + out_cand = v4_attention_v1( toy["q"], toy["k"], toy["v"], @@ -357,7 +357,7 @@ def test_g25_determinism_fp32_mha(): ) scale = 1.0 / math.sqrt(toy["D"]) - out_a = v4_attention( + out_a = v4_attention_v1( toy["q"], toy["k"], toy["v"], @@ -368,7 +368,7 @@ def test_g25_determinism_fp32_mha(): training=False, scale=scale, ) - out_b = v4_attention( + out_b = v4_attention_v1( toy["q"], toy["k"], toy["v"], @@ -379,7 +379,7 @@ def test_g25_determinism_fp32_mha(): training=False, scale=scale, ) - assert torch.equal(out_a, out_b), "v4_attention FWD is non-deterministic at fp32 / MHA" + 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(): @@ -396,7 +396,7 @@ def test_g25_dropout_with_training_is_rejected(): ) scale = 1.0 / math.sqrt(toy["D"]) with pytest.raises(NotImplementedError, match="dropout"): - v4_attention( + v4_attention_v1( toy["q"], toy["k"], toy["v"], diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p26_dispatch.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p26_dispatch.py index 5dacfc7ac..fcc632f2c 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p26_dispatch.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p26_dispatch.py @@ -4,30 +4,20 @@ # See LICENSE for license information. ############################################################################### -"""Plan-4 P26 — dispatch precedence in :class:`DeepseekV4Attention` (CSA path). +"""CSA (cr=4) attention dispatch plumbing in :class:`DeepseekV4Attention`. -Asserts the static plumbing that wires ``use_v4_triton_csa_attention`` -into the V4 CSA forward path, without depending on a full forward -pass: +Static (wiring) gate for the unified ``use_v4_csa_attention_backend`` selector: -* :class:`DeepSeekV4TransformerConfig` already exposes - ``use_v4_triton_csa_attention: bool = False`` (P25 added it; P26 wires - it through ``DeepseekV4Attention.__init__`` + ``_csa_forward``). -* :class:`DeepseekV4Attention.__init__` reads - ``config.use_v4_triton_csa_attention`` and stores it as - ``self._use_v4_triton_csa_attention``, -* the flag is auto-disabled when ``compress_ratio != 4`` (cr ∈ {0, 128} - layers use the separate ``use_v4_triton_attention`` flag landing in - P25), -* :meth:`DeepseekV4Attention._csa_forward` reads - ``self._use_v4_triton_csa_attention`` and dispatches to - :func:`v4_csa_attention` when set, else - :func:`eager_v4_csa_attention`. +* :class:`DeepSeekV4TransformerConfig` exposes ``use_v4_csa_attention_backend: str`` + defaulting to ``"triton_v1"``; +* :class:`DeepseekV4Attention.__init__` reads it into ``self._csa_backend``; +* :meth:`DeepseekV4Attention._csa_forward` switches on ``self._csa_backend``, + dispatching to ``v4_csa_attention_gluon`` / ``v4_csa_attention_v2`` / + ``v4_csa_attention_v1`` (pool) / ``v4_csa_attention_v0`` (gathered, deprecated) + / ``eager_v4_csa_attention``. -The full kernel-vs-eager output equivalence is covered by G26 / G27 -(``test_v4_p26_v4_csa_attention_fwd.py`` / -``test_v4_p26_v4_csa_attention_bwd.py``); this file is the static / -wiring-side gate. +Output equivalence is covered by the G26/G27 fwd/bwd tests; this is the static +wiring gate. """ from __future__ import annotations @@ -36,87 +26,44 @@ DeepSeekV4TransformerConfig, ) -# --------------------------------------------------------------------------- -# Static plumbing — config + module surface -# --------------------------------------------------------------------------- +def test_config_exposes_csa_backend_selector(): + """V4 config exposes ``use_v4_csa_attention_backend: str = 'triton_v1'``.""" + fields = DeepSeekV4TransformerConfig.__dataclass_fields__ + assert "use_v4_csa_attention_backend" in fields + assert fields["use_v4_csa_attention_backend"].default == "triton_v1" -def test_p26_config_exposes_use_v4_triton_csa_attention(): - """V4 transformer config exposes ``use_v4_triton_csa_attention: bool = False``.""" - fields = {f.name for f in DeepSeekV4TransformerConfig.__dataclass_fields__.values()} - assert "use_v4_triton_csa_attention" in fields - csa_field = DeepSeekV4TransformerConfig.__dataclass_fields__["use_v4_triton_csa_attention"] - # Default must be False so existing checkpoints / smokes are - # unaffected by the new switch. - assert csa_field.default is False - - -def test_p26_attention_init_reads_use_v4_triton_csa_attention_flag(): - """``DeepseekV4Attention.__init__`` reads the CSA config flag once.""" - from primus.backends.megatron.core.transformer.deepseek_v4_attention import ( - DeepseekV4Attention, - ) - - init_consts = DeepseekV4Attention.__init__.__code__.co_consts - assert "use_v4_triton_csa_attention" in init_consts, ( - "DeepseekV4Attention.__init__ must read config.use_v4_triton_csa_attention " "(plan-4 P26 plumbing)." - ) - - -def test_p26_attention_init_auto_disables_for_non_csa_layers(): - """``__init__`` auto-disables the CSA flag for cr != 4 layers. - - The dispatch contract is: - cr ∈ {0, 128} ← ``use_v4_triton_attention`` (P25) - cr == 4 ← ``use_v4_triton_csa_attention`` (P26) - - The init's auto-disable branch surfaces this contract explicitly so - a stray run script with ``use_v4_triton_csa_attention=True`` does - not silently accelerate the cr == 4 layers only and skew - apples-to-apples perf comparisons. - """ +def test_attention_init_reads_csa_backend_selector(): + """``__init__`` reads config.use_v4_csa_attention_backend into ``_csa_backend``.""" from primus.backends.megatron.core.transformer.deepseek_v4_attention import ( DeepseekV4Attention, ) - # The init body references the flag name as a string constant via - # ``getattr(config, "use_v4_triton_csa_attention", False)`` and - # checks ``self.compress_ratio != 4`` to flip it back off. - init_consts = DeepseekV4Attention.__init__.__code__.co_consts - assert "use_v4_triton_csa_attention" in init_consts - # The integer constant 4 is the constexpr that gates the - # auto-disable branch (matches plan-4 02-phase-details Phase 26 - # design notes). - assert 4 in init_consts + init_consts = " ".join(str(c) for c in DeepseekV4Attention.__init__.__code__.co_consts if c is not None) + assert "use_v4_csa_attention_backend" in init_consts -def test_p26_csa_forward_consults_use_v4_triton_csa_attention_flag(): - """``_csa_forward`` reads ``self._use_v4_triton_csa_attention`` and references the kernel.""" +def test_csa_forward_switches_on_csa_backend_and_references_kernels(): + """``_csa_forward`` switches on ``self._csa_backend`` and references each entry.""" from primus.backends.megatron.core.transformer.deepseek_v4_attention import ( DeepseekV4Attention, ) co = DeepseekV4Attention._csa_forward.__code__ names = set(co.co_names) - consts_str = " ".join(str(c) for c in co.co_consts if c is not None) - - # The runtime flag attribute must be referenced from the bytecode. - # Python compiles ``self._use_v4_triton_csa_attention`` so the - # attribute name lives in ``co_names``; check both for safety. - has_flag_ref = "_use_v4_triton_csa_attention" in names or any( - "_use_v4_triton_csa_attention" in str(c) for c in co.co_consts - ) - assert has_flag_ref, ( - "DeepseekV4Attention._csa_forward must consult " - "self._use_v4_triton_csa_attention (plan-4 P26 dispatch — precedence " - "use_v4_triton_csa_attention > eager)." - ) - - # The Triton kernel function must be referenced (P31 routes through - # the pool/topk variant to avoid materialising gathered tensors). - refs_csa_kernel = "v4_csa_attention_from_pool" in names or "v4_csa_attention_from_pool" in consts_str - assert refs_csa_kernel, ( - "DeepseekV4Attention._csa_forward must call v4_csa_attention_from_pool " - "when the CSA Triton flag is on (plan-5 P31)." - ) + consts = " ".join(str(c) for c in co.co_consts if c is not None) + + assert "_csa_backend" in names + for be in ("gluon", "triton_v2", "triton_v1", "flydsl_v0"): + assert be in consts, f"_csa_forward must handle CSA backend value {be!r}" + # references the pool (v1), sparse-MLA (v2), gathered (v0) and eager entries; + # gluon is loaded lazily onto ``self._v4_csa_attention_gluon`` (attribute access). + for fn in ( + "v4_csa_attention_v1", + "v4_csa_attention_v2", + "_v4_csa_attention_gluon", + "v4_csa_attention_v0", + "eager_v4_csa_attention", + ): + assert fn in names or fn in consts, f"_csa_forward must reference {fn}" diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p26_v4_csa_attention_bwd.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p26_v4_csa_attention_bwd.py index 4d3ea2804..fedbf818f 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p26_v4_csa_attention_bwd.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p26_v4_csa_attention_bwd.py @@ -4,10 +4,10 @@ # See LICENSE for license information. ############################################################################### -"""Plan-4 P26 G27 — `v4_csa_attention` Triton BWD equivalence to autograd-on-eager. +"""Plan-4 P26 G27 — `v4_csa_attention_v0` Triton BWD equivalence to autograd-on-eager. Asserts that gradients (``dq``, ``dk_local``, ``dv_local``, -``dgathered``, ``dsink``) returned by :func:`v4_csa_attention`'s +``dgathered``, ``dsink``) returned by :func:`v4_csa_attention_v0`'s autograd Function match the gradients computed by autograd-on- :func:`eager_v4_csa_attention` within the plan-4 tolerance budget. @@ -36,13 +36,13 @@ torch = pytest.importorskip("torch") if not torch.cuda.is_available(): - pytest.skip("v4_csa_attention Triton kernel requires CUDA / HIP", allow_module_level=True) + 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, + v4_csa_attention_v0, ) # --------------------------------------------------------------------------- @@ -256,7 +256,7 @@ def test_g27_csa_bwd_matches_eager( ref_inp["sink"], ) - out_cand = v4_csa_attention( + out_cand = v4_csa_attention_v0( cand_inp["q"], cand_inp["k_local"], cand_inp["v_local"], @@ -326,7 +326,7 @@ def test_g27_csa_bwd_grads_finite_with_masked_topk(dtype: torch.dtype): ) sink = torch.randn(H, generator=g, device=device, dtype=torch.float32, requires_grad=True) * 0.1 - out = v4_csa_attention( + out = v4_csa_attention_v0( q, k_local, v_local, diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p26_v4_csa_attention_fwd.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p26_v4_csa_attention_fwd.py index a7f42414c..73383ac29 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p26_v4_csa_attention_fwd.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p26_v4_csa_attention_fwd.py @@ -4,9 +4,9 @@ # See LICENSE for license information. ############################################################################### -"""Plan-4 P26 G26 — `v4_csa_attention` Triton FWD equivalence to eager. +"""Plan-4 P26 G26 — `v4_csa_attention_v0` Triton FWD equivalence to eager. -Asserts that :func:`v4_csa_attention` (Triton kernel from +Asserts that :func:`v4_csa_attention_v0` (Triton kernel from ``primus...transformer.v4_attention_kernels.v4_csa_attention``) produces forward output equal to :func:`eager_v4_csa_attention` within the plan-4 tolerance budget across: @@ -18,7 +18,7 @@ * fp32 and bf16 inputs; * ``sink ∈ {None, learned [H]}``; * ``K_topk == 0`` short-circuit — the wrapper falls through to the - dense :func:`v4_attention` kernel and the result must still match + dense :func:`v4_attention_v1` kernel and the result must still match the eager reference's degenerate (``gathered.shape[2] == 0``) limit. The test is GPU-only (Triton requires CUDA / HIP); CPU runs are @@ -34,13 +34,13 @@ torch = pytest.importorskip("torch") if not torch.cuda.is_available(): - pytest.skip("v4_csa_attention Triton kernel requires CUDA / HIP", allow_module_level=True) + 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, + v4_csa_attention_v0, ) # --------------------------------------------------------------------------- @@ -223,7 +223,7 @@ def test_g26_csa_fwd_matches_eager( scale=scale, ) - out_cand = v4_csa_attention( + out_cand = v4_csa_attention_v0( toy["q"], toy["k_local"], toy["v_local"], @@ -253,7 +253,7 @@ def test_g26_csa_fwd_matches_eager( @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_g26_csa_fwd_short_circuits_when_k_topk_zero(dtype: torch.dtype, sink_on: bool): - """``v4_csa_attention(..., gathered.shape[2] == 0, ...)`` falls through to dense.""" + """``v4_csa_attention_v0(..., gathered.shape[2] == 0, ...)`` falls through to dense.""" B, H, S, D = 1, 4, 32, 64 swa_window = 16 @@ -285,8 +285,8 @@ def test_g26_csa_fwd_short_circuits_when_k_topk_zero(dtype: torch.dtype, sink_on scale=scale, ) - # Candidate: wrapper short-circuits to dense v4_attention. - out_cand = v4_csa_attention( + # Candidate: wrapper short-circuits to dense v4_attention_v1. + out_cand = v4_csa_attention_v0( q, k_local, v_local, @@ -309,7 +309,7 @@ def test_g26_csa_fwd_short_circuits_when_k_topk_zero(dtype: torch.dtype, sink_on def test_g26_csa_dropout_with_training_is_rejected(): - """``v4_csa_attention`` raises NotImplementedError on dropout + training=True. + """``v4_csa_attention_v0`` raises NotImplementedError on dropout + training=True. V4 trains with attn_dropout=0; the kernel does not implement in-kernel attention dropout. We refuse explicitly so a stray @@ -329,7 +329,7 @@ def test_g26_csa_dropout_with_training_is_rejected(): sparse_mask = torch.zeros(B, S, K_topk, device=device, dtype=dtype) with pytest.raises(NotImplementedError, match="does not implement in-kernel"): - v4_csa_attention( + v4_csa_attention_v0( q, k_local, v_local, diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p27_dispatch_precedence.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p27_dispatch_precedence.py index 876254a0b..09273dd04 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p27_dispatch_precedence.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p27_dispatch_precedence.py @@ -4,72 +4,25 @@ # See LICENSE for license information. ############################################################################### -"""Plan-4 P27 G29 — runtime dispatch precedence in :class:`DeepseekV4Attention`. - -The static plumbing tests -(``test_v4_p25_dispatch.py`` and ``test_v4_p26_dispatch.py``) cover the -*config-side* / *bytecode-side* wiring: that the flags exist, are -read by ``__init__``, are auto-disabled for the wrong layer kind, and -that ``forward`` / ``_csa_forward`` reference the runtime flag -attributes. - -This file extends them with **runtime mock-based** assertions — for -each ``(compress_ratio, flags) → expected kernel`` combination -documented in the :class:`DeepseekV4Attention` class docstring, we -mock all four candidate kernels and assert exactly one fires: - - cr == 0: - use_turbo_attention=T → core_attention (Plan-3 P22) - use_turbo_attention=F, use_v4_triton_attention=T → v4_attention (P25) - use_turbo_attention=F, use_v4_triton_attention=F → eager_v4_attention (P24) - - cr == 128: - use_v4_triton_attention=T → v4_attention (HCA path) (P25) - use_v4_triton_attention=F → eager_v4_attention (P24) - - cr == 4: - use_v4_triton_csa_attention=T → v4_csa_attention_from_pool (P31) - use_v4_triton_csa_attention=F → eager_v4_csa_attention (P24) - -We bypass ``__init__`` (which requires distributed + GPU + Megatron -config) by allocating the attention instance with ``__new__`` and -populating only the attributes the dispatch path reads. This keeps -the test CPU-only and instantaneous while still exercising the actual -``forward`` / ``_csa_forward`` Python code paths from the production -module. - -The kernel-vs-eager numerical equivalence is covered by G23/G24 (P25) -and G26/G27 (P26); G28 (P27) extends those to release-tier shapes. -G29 is the *dispatch* gate — it ensures the right kernel is selected -for each valid configuration, which the static tests cannot prove -without running the forward path. - -Also covers :meth:`DeepseekV4Attention._log_kernel_choice` — the new -P27 ``INFO`` log line summarising each layer's kernel choice — by -asserting the expected string content for each of the 7 valid -(cr, flag) combinations. +"""Runtime dispatch of :class:`DeepseekV4Attention` on the unified selectors. + +Mock-based gate for the ``use_v4_attention_backend`` (dense/HCA) and +``use_v4_csa_attention_backend`` (CSA) string selectors. For each +``(compress_ratio, backend) -> expected kernel`` combination we patch the +candidate kernels and assert exactly one fires. ``__init__`` is bypassed via +``__new__`` (CPU-only) so we exercise the real ``forward`` / +``_attention_backend_forward`` / ``_csa_forward`` / ``_log_kernel_choice`` code. """ from __future__ import annotations +import contextlib import logging from unittest.mock import MagicMock, patch import pytest # isort: off -# Import order matters here: the V4 transformer config MUST be imported -# FIRST so the ``primus...models.deepseek_v4`` package's ``__init__`` -# chain (which in turn imports ``deepseek_v4_block`` and -# ``deepseek_v4_attention``) finishes wiring up before we grab the -# attention module by reference. Without this priming step the -# subsequent ``import ... deepseek_v4_attention`` races against -# ``deepseek_v4_block.py`` re-importing ``DeepseekV4Attention`` from -# the partially-initialised module and raises ``ImportError``. The -# P25 / P26 dispatch tests dodge the same race by deferring the import -# to inside each test function; we want a module-level handle so the -# test fixtures (caplog parametrisations) can reference it directly, -# so we hard-pin the order with an ``isort: off`` block. from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( # noqa: F401 DeepSeekV4TransformerConfig, ) @@ -80,139 +33,142 @@ # isort: on -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - def _make_bare_attn( *, - compress_ratio: int, - use_core_attention: bool = False, - use_v4_triton_attention: bool = False, - use_v4_triton_csa_attention: bool = False, - use_v4_tilelang_attention: bool = False, - use_v4_tilelang_csa_attention: bool = False, - layer_number: int = 0, -) -> DeepseekV4Attention: - """Build a bare :class:`DeepseekV4Attention` for runtime dispatch tests. - - Bypasses ``__init__`` (which requires Megatron / GPU / distributed) - and populates only the attributes the dispatch path consults. The - helper methods on the real class read these attributes by name; we - do not need to subclass. - - Plan-8 P57 close-out 2 added the two ``_use_v4_tilelang_*`` flags; - they default False so existing tests remain on the Triton path. - """ + compress_ratio, + attn_backend="triton_v1", + csa_backend="triton_v1", + use_core_attention=False, + layer_number=0, +): attn = DeepseekV4Attention.__new__(DeepseekV4Attention) attn.compress_ratio = int(compress_ratio) attn.layer_number = int(layer_number) attn._use_core_attention = bool(use_core_attention) - attn._use_v4_triton_attention = bool(use_v4_triton_attention) - attn._use_v4_triton_csa_attention = bool(use_v4_triton_csa_attention) - attn._use_v4_tilelang_attention = bool(use_v4_tilelang_attention) - attn._use_v4_tilelang_csa_attention = bool(use_v4_tilelang_csa_attention) + attn._attn_backend = attn_backend + attn._csa_backend = csa_backend + # gluon is loaded lazily onto the instance in __init__; mirror the None + # placeholders here so non-gluon paths never touch it (tests that exercise the + # gluon branch set these to a mock explicitly). + attn._v4_attention_gluon = None + attn._v4_csa_attention_gluon = None return attn # --------------------------------------------------------------------------- -# G29.A — _log_kernel_choice emits the correct kernel name per (cr, flags) +# _log_kernel_choice # --------------------------------------------------------------------------- @pytest.mark.parametrize( - "compress_ratio,use_core,use_v4t,use_csa,expected_substr", + "cr,use_core,attn_be,csa_be,expected", [ - # cr == 0: 3 valid kernel choices - pytest.param(0, True, False, False, "core_attention", id="cr0_turbo"), - pytest.param(0, False, True, False, "Triton, dense path", id="cr0_triton"), - pytest.param(0, False, False, False, "eager Python, dense path", id="cr0_eager"), - # cr == 128: 2 valid kernel choices (no turbo path for HCA) - pytest.param(128, False, True, False, "Triton, HCA path", id="cr128_triton"), - pytest.param(128, False, False, False, "eager Python, HCA path", id="cr128_eager"), - # cr == 4: 2 valid kernel choices (no turbo / dense-triton path for CSA) - pytest.param(4, False, False, True, "v4_csa_attention_from_pool (Triton)", id="cr4_triton"), - pytest.param(4, False, False, False, "v4_csa_attention (eager Python)", id="cr4_eager"), + pytest.param(0, True, "triton_v1", "triton_v1", "core_attention", id="cr0_turbo"), + pytest.param(0, False, "gluon", "triton_v1", "dense attention backend = gluon", id="cr0_gluon"), + pytest.param( + 0, False, "triton_v1", "triton_v1", "dense attention backend = triton_v1", id="cr0_triton_v1" + ), + pytest.param( + 128, False, "triton_v2", "triton_v1", "HCA attention backend = triton_v2", id="cr128_v2" + ), + pytest.param(4, False, "triton_v1", "gluon", "CSA attention backend = gluon", id="cr4_gluon"), + pytest.param(4, False, "triton_v1", "eager", "CSA attention backend = eager", id="cr4_eager"), ], ) -def test_p27_log_kernel_choice_emits_expected_kernel( - compress_ratio: int, - use_core: bool, - use_v4t: bool, - use_csa: bool, - expected_substr: str, - caplog: pytest.LogCaptureFixture, -) -> None: - """:meth:`_log_kernel_choice` emits one INFO line per layer naming the kernel. - - Plan-4 P27 surfaces the dispatch outcome at boot via - :meth:`DeepseekV4Attention._log_kernel_choice` so smoke / - training logs unambiguously show which kernel each layer is - firing through. The log line format is checked by - :func:`test_p27_log_kernel_choice_format`; this test only - verifies the kernel-name substring is correct. - """ +def test_log_kernel_choice(cr, use_core, attn_be, csa_be, expected, caplog): attn = _make_bare_attn( - compress_ratio=compress_ratio, + compress_ratio=cr, + attn_backend=attn_be, + csa_backend=csa_be, use_core_attention=use_core, - use_v4_triton_attention=use_v4t, - use_v4_triton_csa_attention=use_csa, layer_number=17, ) caplog.set_level(logging.INFO, logger=v4_attn_mod.__name__) attn._log_kernel_choice() - assert expected_substr in caplog.text, ( - f"_log_kernel_choice did not emit the expected kernel name for cr={compress_ratio} " - f"(use_core={use_core}, use_v4t={use_v4t}, use_csa={use_csa}); " - f"expected substring '{expected_substr}', got log text:\n{caplog.text}" - ) + assert expected in caplog.text, f"expected {expected!r} in log, got:\n{caplog.text}" -def test_p27_log_kernel_choice_format(caplog: pytest.LogCaptureFixture) -> None: - """Log line format is ``[V4-attn] Layer N: cr=R, kernel = ...``.""" +def test_log_kernel_choice_format(caplog): attn = _make_bare_attn(compress_ratio=0, layer_number=42) caplog.set_level(logging.INFO, logger=v4_attn_mod.__name__) attn._log_kernel_choice() - assert "[V4-attn]" in caplog.text - assert "Layer 42" in caplog.text - assert "cr=0" in caplog.text - assert "kernel = " in caplog.text - - -def test_p27_log_kernel_choice_emits_once_per_call(caplog: pytest.LogCaptureFixture) -> None: - """A single :meth:`_log_kernel_choice` call emits exactly one INFO line.""" - attn = _make_bare_attn(compress_ratio=4, use_v4_triton_csa_attention=True, layer_number=3) - caplog.set_level(logging.INFO, logger=v4_attn_mod.__name__) - attn._log_kernel_choice() - v4_attn_records = [r for r in caplog.records if r.name == v4_attn_mod.__name__] - assert len(v4_attn_records) == 1 + assert "[V4-attn]" in caplog.text and "Layer 42" in caplog.text + assert "cr=0" in caplog.text and "kernel = " in caplog.text # --------------------------------------------------------------------------- -# G29.B — runtime mock dispatch on ``_attention_forward_via_v4_triton`` +# dense/HCA dispatch: _attention_backend_forward # --------------------------------------------------------------------------- -def test_p27_via_v4_triton_helper_calls_v4_attention_kernel() -> None: - """:meth:`_attention_forward_via_v4_triton` calls the Triton kernel. +@pytest.mark.parametrize( + "attn_be,fired", + [ + ("gluon", "v4_attention_gluon"), + ("triton_v2", "v4_attention_v2"), + ], +) +def test_dense_backend_forward_dispatch(attn_be, fired): + attn = _make_bare_attn(compress_ratio=0, attn_backend=attn_be) + attn.attn_sink = None + attn.attn_sliding_window = 128 + attn.attn_dropout = 0.0 + attn.training = False + attn._attention_scale = MagicMock(return_value=0.125) + sentinel = object() + mock_fired = MagicMock(return_value=sentinel) + # gluon is dispatched via the lazily-loaded instance attribute + # (self._v4_attention_gluon); other backends via a module-level entry. + if attn_be == "gluon": + attn._v4_attention_gluon = mock_fired + ctx = contextlib.nullcontext() + else: + ctx = patch.object(v4_attn_mod, fired, mock_fired) + with ctx: + out = DeepseekV4Attention._attention_backend_forward( + attn, + MagicMock(), + MagicMock(), + MagicMock(), + additive_mask=None, + hca_local_seqlen=0, + S=8, + device="cpu", + dtype=None, + ) + assert out is sentinel + mock_fired.assert_called_once() + + +def test_dense_backend_forward_triton_v1_calls_launcher_helper(): + attn = _make_bare_attn(compress_ratio=0, attn_backend="triton_v1") + attn.attn_sliding_window = 128 + attn._attention_forward_via_v4_triton = MagicMock(return_value="TRITONV1") + out = DeepseekV4Attention._attention_backend_forward( + attn, + MagicMock(), + MagicMock(), + MagicMock(), + additive_mask=None, + hca_local_seqlen=0, + S=8, + device="cpu", + dtype=None, + ) + assert out == "TRITONV1" + attn._attention_forward_via_v4_triton.assert_called_once() + - P25 routes the dense / HCA Triton path through this helper. The - static test (P25) checks the bytecode references the kernel name; - this runtime test patches the kernel in the module namespace and - asserts it is actually invoked. Combined with the dispatch tests - below, this confirms the Triton kernel is the one that fires - when ``use_v4_triton_attention=True``. - """ - attn = _make_bare_attn(compress_ratio=0, use_v4_triton_attention=True) +def test_via_v4_triton_helper_calls_v4_attention_v1(): + attn = _make_bare_attn(compress_ratio=0, attn_backend="triton_v1") attn.attn_sink = None attn.attn_sliding_window = 0 attn.attn_dropout = 0.0 attn.training = False attn._attention_scale = MagicMock(return_value=0.125) - sentinel = object() - with patch.object(v4_attn_mod, "v4_attention", return_value=sentinel) as mock_kernel: + with patch.object(v4_attn_mod, "v4_attention_v1", return_value=sentinel) as mock_kernel: out = DeepseekV4Attention._attention_forward_via_v4_triton( attn, q=MagicMock(), @@ -220,217 +176,67 @@ def test_p27_via_v4_triton_helper_calls_v4_attention_kernel() -> None: v=MagicMock(), attn_mask=MagicMock(), ) - assert out is sentinel mock_kernel.assert_called_once() - kwargs = mock_kernel.call_args.kwargs - assert "additive_mask" in kwargs - assert kwargs["swa_window"] == 0 # --------------------------------------------------------------------------- -# G29.C — runtime mock dispatch on ``_csa_forward`` +# CSA dispatch: _csa_forward # --------------------------------------------------------------------------- def _stub_csa_inputs(): - """Return MagicMock tensor stubs accepted by ``_csa_forward``. - - Each value supports the few attribute / method accesses - ``_csa_forward`` makes before reaching the dispatch line. We - intercept the prep ops via ``patch.object`` so the stub tensors - are never actually arithmetised. - """ - hidden = MagicMock() q_bh = MagicMock() q_bh.shape = (1, 2, 3, 4) # B, H, S, Dh - k_local_bh = MagicMock() - v_local_bh = MagicMock() - local_mask = MagicMock() - return hidden, q_bh, k_local_bh, v_local_bh, local_mask + return MagicMock(), q_bh, MagicMock(), MagicMock(), MagicMock() @pytest.mark.parametrize( - "use_csa,fired_attr,suppressed_attr", + "csa_be,fired,builds_gathered", [ - pytest.param( - True, - "v4_csa_attention_from_pool", - "eager_v4_csa_attention", - id="csa_triton_wins", - ), - pytest.param( - False, - "eager_v4_csa_attention", - "v4_csa_attention_from_pool", - id="csa_eager_when_off", - ), + ("gluon", "v4_csa_attention_gluon", False), + ("triton_v2", "v4_csa_attention_v2", False), + ("triton_v1", "v4_csa_attention_v1", False), + ("eager", "eager_v4_csa_attention", True), ], ) -def test_p27_csa_forward_dispatch_precedence( - use_csa: bool, - fired_attr: str, - suppressed_attr: str, -) -> None: - """``_csa_forward`` picks v4_csa_attention_from_pool vs eager_v4_csa_attention by flag. - - Precedence (P26): ``use_v4_triton_csa_attention > eager``. - """ - attn = _make_bare_attn( - compress_ratio=4, - use_v4_triton_csa_attention=use_csa, - ) +def test_csa_forward_dispatch(csa_be, fired, builds_gathered): + attn = _make_bare_attn(compress_ratio=4, csa_backend=csa_be) attn.attn_sink = None - attn.attn_sliding_window = 0 + attn.attn_sliding_window = 128 attn.attn_dropout = 0.0 attn.training = False attn._attention_scale = MagicMock(return_value=0.125) - # Patch the prep ops invoked before the kernel dispatch. pool_mock = MagicMock() pool_mock.shape = (1, 8, 4) # [B, P, Dh] pool_mock.unsqueeze.return_value.expand.return_value = MagicMock() attn._build_compressed_pool = MagicMock(return_value=pool_mock) - indexer_topk = MagicMock() - indexer_topk.shape = (1, 3, 4) # [B, S, K] - indexer_topk.clamp.return_value = MagicMock() - indexer_topk.__ge__ = MagicMock(return_value=MagicMock()) - attn.indexer = MagicMock(return_value=(indexer_topk, MagicMock())) - - fired_sentinel = object() - suppressed_called = [False] - - def _suppressed(*args, **kwargs): - suppressed_called[0] = True - return object() + topk = MagicMock() + topk.shape = (1, 3, 4) # [B, S, K] + topk.clamp.return_value = MagicMock() + topk.__ge__ = MagicMock(return_value=MagicMock()) + attn.indexer = MagicMock(return_value=(topk, MagicMock())) - with patch.object(v4_attn_mod, fired_attr, return_value=fired_sentinel) as mock_fired, patch.object( - v4_attn_mod, suppressed_attr, side_effect=_suppressed - ), patch("torch.gather", return_value=MagicMock()), patch("torch.where", return_value=MagicMock()): - hidden, q_bh, k_local_bh, v_local_bh, local_mask = _stub_csa_inputs() + sentinel = object() + mock_fired = MagicMock(return_value=sentinel) + # gluon is dispatched via the lazily-loaded instance attribute + # (self._v4_csa_attention_gluon); other backends via a module-level entry. + if csa_be == "gluon": + attn._v4_csa_attention_gluon = mock_fired + fired_ctx = contextlib.nullcontext() + else: + fired_ctx = patch.object(v4_attn_mod, fired, mock_fired) + with fired_ctx, patch("torch.gather", return_value=MagicMock()), patch( + "torch.where", return_value=MagicMock() + ): + _, q_bh, k_local_bh, v_local_bh, local_mask = _stub_csa_inputs() out = DeepseekV4Attention._csa_forward( attn, - hidden=hidden, + hidden=MagicMock(), q_bh=q_bh, k_local_bh=k_local_bh, v_local_bh=v_local_bh, local_mask=local_mask, ) - - assert out is fired_sentinel, ( - f"_csa_forward did not return the expected kernel's output for " - f"use_v4_triton_csa_attention={use_csa}; expected {fired_attr} to fire." - ) + assert out is sentinel mock_fired.assert_called_once() - assert not suppressed_called[0], ( - f"Both kernel paths were invoked when use_v4_triton_csa_attention={use_csa}; " - f"the suppressed kernel '{suppressed_attr}' must NOT be called." - ) - - -def test_p27_csa_forward_passes_kernel_args_through() -> None: - """:meth:`_csa_forward` passes the prepared kwargs to the chosen kernel. - - This is a contract check on the kwargs the kernel sees: the same - set of keyword arguments must reach ``v4_csa_attention`` and - ``eager_v4_csa_attention``, otherwise a kernel swap would silently - re-route through a different code path. - """ - common_expected_keys = { - "sink", - "swa_window", - "attn_dropout", - "training", - "scale", - } - for use_csa, target in [(True, "v4_csa_attention_from_pool"), (False, "eager_v4_csa_attention")]: - attn = _make_bare_attn(compress_ratio=4, use_v4_triton_csa_attention=use_csa) - attn.attn_sink = MagicMock() - attn.attn_sliding_window = 128 - attn.attn_dropout = 0.0 - attn.training = False - attn._attention_scale = MagicMock(return_value=0.125) - pool_mock = MagicMock() - pool_mock.shape = (1, 8, 4) - pool_mock.unsqueeze.return_value.expand.return_value = MagicMock() - attn._build_compressed_pool = MagicMock(return_value=pool_mock) - indexer_topk = MagicMock() - indexer_topk.shape = (1, 3, 4) - indexer_topk.clamp.return_value = MagicMock() - indexer_topk.__ge__ = MagicMock(return_value=MagicMock()) - attn.indexer = MagicMock(return_value=(indexer_topk, MagicMock())) - with patch.object(v4_attn_mod, target, return_value=MagicMock()) as mock_kernel, patch( - "torch.gather", return_value=MagicMock() - ), patch("torch.where", return_value=MagicMock()): - hidden, q_bh, k_local_bh, v_local_bh, local_mask = _stub_csa_inputs() - DeepseekV4Attention._csa_forward( - attn, - hidden=hidden, - q_bh=q_bh, - k_local_bh=k_local_bh, - v_local_bh=v_local_bh, - local_mask=local_mask, - ) - kwargs = mock_kernel.call_args.kwargs - expected_keys = set(common_expected_keys) - expected_keys.add("topk_idxs" if use_csa else "sparse_mask") - missing = expected_keys - set(kwargs.keys()) - assert not missing, ( - f"_csa_forward did not pass {sorted(missing)} to {target} when " - f"use_v4_triton_csa_attention={use_csa}; check kwargs={list(kwargs.keys())}." - ) - - -# --------------------------------------------------------------------------- -# G29.D — turbo precedence dominates v4_triton on dense (cr == 0) -# --------------------------------------------------------------------------- - - -def test_p27_init_auto_disables_v4_triton_for_csa_layers() -> None: - """``__init__`` must auto-disable ``use_v4_triton_attention`` for cr == 4. - - The contract is documented in the class docstring under - "Auto-disable rules". The static test (P25) verifies the - bytecode contains the (0, 128) constexpr; this test verifies - the runtime attribute is correctly flipped on a bare instance. - """ - # Mirror the init body: cr=4 + use_v4_triton_attention=True → flag flipped off. - attn = _make_bare_attn(compress_ratio=4, use_v4_triton_attention=True) - # Manually run the init's auto-disable predicate (this is the - # exact source of the runtime guard). - if attn._use_v4_triton_attention and attn.compress_ratio not in (0, 128): - attn._use_v4_triton_attention = False - assert attn._use_v4_triton_attention is False, ( - "DeepseekV4Attention.__init__ must auto-disable use_v4_triton_attention for " - "compress_ratio=4 layers (CSA opts in via use_v4_triton_csa_attention; plan-4 P25)." - ) - - -def test_p27_init_auto_disables_v4_triton_csa_for_dense_hca_layers() -> None: - """``__init__`` must auto-disable ``use_v4_triton_csa_attention`` for cr != 4.""" - for cr in (0, 128): - attn = _make_bare_attn(compress_ratio=cr, use_v4_triton_csa_attention=True) - if attn._use_v4_triton_csa_attention and attn.compress_ratio != 4: - attn._use_v4_triton_csa_attention = False - assert attn._use_v4_triton_csa_attention is False, ( - f"__init__ must auto-disable use_v4_triton_csa_attention for compress_ratio={cr} " - "layers (dense / HCA opt in via use_v4_triton_attention; plan-4 P26)." - ) - - -def test_p27_log_message_mentions_layer_number() -> None: - """The startup log line includes the layer number for cross-rank correlation. - - Each rank's log file holds entries for the layers it owns; the - explicit ``Layer N`` makes it easy to grep for a specific - layer's kernel choice in distributed smoke / training logs. - """ - layer_numbers = [0, 1, 7, 17, 59] - for ln in layer_numbers: - attn = _make_bare_attn(compress_ratio=128, use_v4_triton_attention=True, layer_number=ln) - with patch.object(v4_attn_mod.logger, "info") as mock_info: - attn._log_kernel_choice() - mock_info.assert_called_once() - # logger.info uses %-style formatting; the layer number is the - # first %-arg, the cr is the second, and the kernel name is the third. - call_args = mock_info.call_args - assert call_args.args[1] == ln, f"Expected layer_number {ln}, got {call_args.args[1]}" diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p31_v4_csa_in_kernel_gather.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p31_v4_csa_in_kernel_gather.py index bc2b0ed43..0f2d8c371 100644 --- a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p31_v4_csa_in_kernel_gather.py +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_p31_v4_csa_in_kernel_gather.py @@ -22,13 +22,13 @@ torch = pytest.importorskip("torch") if not torch.cuda.is_available(): - pytest.skip("v4_csa_attention Triton kernel requires CUDA / HIP", allow_module_level=True) + 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_from_pool, + v4_csa_attention_v1, ) _SHAPES = [ @@ -168,7 +168,7 @@ def test_p31_csa_pool_fwd_matches_gathered_reference( training=False, scale=scale, ) - out_cand = v4_csa_attention_from_pool( + out_cand = v4_csa_attention_v1( inp["q"], inp["k_local"], inp["v_local"], @@ -247,7 +247,7 @@ def test_p31_csa_pool_bwd_matches_gathered_reference( ref_inp["sink"], ) - out_cand = v4_csa_attention_from_pool( + out_cand = v4_csa_attention_v1( cand_inp["q"], cand_inp["k_local"], cand_inp["v_local"], 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..0c972cdd8 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_flydsl_attention.py @@ -0,0 +1,293 @@ +############################################################################### +# 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. +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/third_party/Emerging-Optimizers b/third_party/Emerging-Optimizers index 93d9eb3a6..46eda5a05 160000 --- a/third_party/Emerging-Optimizers +++ b/third_party/Emerging-Optimizers @@ -1 +1 @@ -Subproject commit 93d9eb3a6c899b50de73992826451fba3ab6adfb +Subproject commit 46eda5a0504d321bf45bae278f657ee30ab9a05e diff --git a/tools/ci/check_version_consistency.py b/tools/ci/check_version_consistency.py new file mode 100644 index 000000000..183837011 --- /dev/null +++ b/tools/ci/check_version_consistency.py @@ -0,0 +1,123 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Fail CI on cross-file config drift: a value duplicated across files that fell +out of sync. Runs in the lint job. Covers BASE_IMAGE (ci.yaml vs Dockerfile +ARG), primus.__version__ (PEP 440), pyproject deps vs requirements.txt, action +SHA-pinning, and workflow python-version. Stdlib-only. +""" + +import re +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +CI_YAML = ROOT / ".github/workflows/ci.yaml" +DOCKERFILE = ROOT / ".github/workflows/docker/Dockerfile" +INIT_PY = ROOT / "primus/__init__.py" +PYPROJECT = ROOT / "pyproject.toml" +REQUIREMENTS = ROOT / "requirements.txt" +WORKFLOWS = sorted((ROOT / ".github/workflows").glob("*.y*ml")) + +PEP440 = re.compile(r"^\d+(\.\d+)*((a|b|rc)\d+)?(\.post\d+)?(\.dev\d+)?$") +SHA40 = re.compile(r"^[0-9a-f]{40}$") +NAME_SPEC = re.compile(r"^([A-Za-z0-9._-]+)\s*(.*)$") + + +def _canon(name): + return re.sub(r"[-_.]+", "-", name).lower() + + +def _find(text, pattern, what, where): + match = re.search(pattern, text, re.MULTILINE) + if match is None: + sys.exit(f"ERROR: could not find {what} in {where}") + return match.group(1).strip() + + +def _parse_dep(line): + line = line.split("#", 1)[0].strip() + match = NAME_SPEC.match(line) if line else None + return (_canon(match.group(1)), match.group(2).replace(" ", "")) if match else None + + +def check_base_image(errors): + ci = _find(CI_YAML.read_text(), r"^\s*BASE_IMAGE:\s*(\S+)", "BASE_IMAGE env", CI_YAML) + docker = _find(DOCKERFILE.read_text(), r"^ARG\s+BASE_IMAGE=(\S+)", "ARG BASE_IMAGE default", DOCKERFILE) + if ci != docker: + errors.append(f"BASE_IMAGE drift: ci.yaml={ci!r} vs Dockerfile default={docker!r}.") + + +def check_version(errors): + version = _find(INIT_PY.read_text(), r'__version__\s*=\s*["\']([^"\']+)["\']', "__version__", INIT_PY) + if not PEP440.match(version): + errors.append(f"primus.__version__={version!r} is not a valid PEP 440 version.") + + +def check_deps(errors): + block = re.search(r"^dependencies\s*=\s*\[(.*?)\]", PYPROJECT.read_text(), re.S | re.M) + if block is None: + errors.append("could not find [project].dependencies in pyproject.toml") + return + pyproject_deps = dict(filter(None, (_parse_dep(d) for d in re.findall(r'"([^"]+)"', block.group(1))))) + req = dict(filter(None, (_parse_dep(line) for line in REQUIREMENTS.read_text().splitlines()))) + # pyproject deps must be a subset of requirements with matching specifiers; + # requirements may carry extra dev/CI-only entries (e.g. hip-python). + for name, spec in sorted(pyproject_deps.items()): + if name not in req: + errors.append(f"{name!r} is in pyproject.toml but missing from requirements.txt.") + elif req[name] != spec: + errors.append(f"spec drift for {name!r}: pyproject={spec!r} vs requirements={req[name]!r}.") + + +def check_pinned_actions(errors): + for wf in WORKFLOWS: + for raw in re.findall(r"uses:\s*(\S+)", wf.read_text()): + uses = raw.strip().strip('"').strip("'") + if uses.startswith("./"): + continue + if "@" not in uses: + errors.append(f"{wf.name}: action {uses!r} is not pinned (no @ref).") + elif not SHA40.match(uses.rsplit("@", 1)[1]): + errors.append(f"{wf.name}: action {uses!r} is not pinned to a 40-hex commit SHA.") + + +def check_python_versions(errors): + requires = _find( + PYPROJECT.read_text(), r'requires-python\s*=\s*["\']([^"\']+)["\']', "requires-python", PYPROJECT + ) + floor = re.search(r">=\s*(\d+)\.(\d+)", requires) + min_ver = (int(floor.group(1)), int(floor.group(2))) if floor else None + versions = set() + for wf in WORKFLOWS: + versions.update(re.findall(r'python-version:\s*\[?\s*["\'](\d+\.\d+)["\']', wf.read_text())) + if len(versions) > 1: + errors.append(f"workflow python-version values disagree: {sorted(versions)}.") + for v in versions if min_ver else []: + if tuple(int(x) for x in v.split(".")) < min_ver: + errors.append( + f"workflow python-version {v} is below requires-python >={min_ver[0]}.{min_ver[1]}." + ) + + +def main(): + errors = [] + check_base_image(errors) + check_version(errors) + check_deps(errors) + check_pinned_actions(errors) + check_python_versions(errors) + if errors: + print("CI consistency check FAILED:") + for err in errors: + print(f" - {err}") + return 1 + print("CI consistency OK (base image, version, deps, action pinning, python-version).") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/ci/coverage_summary.py b/tools/ci/coverage_summary.py new file mode 100644 index 000000000..a7477ace8 --- /dev/null +++ b/tools/ci/coverage_summary.py @@ -0,0 +1,175 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Render coverage.py JSON as a compact Markdown table for the CI run summary. + +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). +""" + +import json +import sys +from collections import defaultdict + +OMIT_MODULES = {"tools", "platforms", "pretrain.py"} +# Top-level groups whose sub-packages are shown as indented detail rows; every +# other group (modules, 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. + """ + 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/.py, e.g. pretrain.py + return seg[1], None + if seg[1] in DETAILED_GROUPS: + return seg[1], seg[1] + "/" + seg[2] + return seg[1], None + + +def _pct(covered: int, total: int) -> float: + return (100.0 * covered / total) if total else 0.0 + + +def _aggregate(report: dict): + """Return {group: {detail|group: [covered, statements]}} for kept modules.""" + agg = defaultdict(lambda: defaultdict(lambda: [0, 0])) + for fpath, info in report.get("files", {}).items(): + result = classify(fpath) + if result is None: + continue + group, detail = result + if group in OMIT_MODULES: + continue + s = info["summary"] + a = agg[group][detail or group] + a[0] += s["covered_lines"] + a[1] += s["num_statements"] + return agg + + +def render(primary: dict, title: str, secondaries: list = None) -> str: + pa = _aggregate(primary) + sas = [_aggregate(s) for s in (secondaries or [])] + two = bool(sas) + + def e2e_cov(group, key): + # Merge E2E reports by max covered lines (jobs cover near-disjoint + # 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): + 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) + " |" + + 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. + 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): + 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)] + + tc = sum(group_totals(g)[0] for g in groups) + tn = sum(group_totals(g)[1] for g in groups) + te = sum(group_totals(g)[2] for g in groups) if two else 0 + excl = ", ".join(sorted(OMIT_MODULES)) + + # coverage's own totals over every primus file (nothing excluded), for context. + p_all = primary.get("totals", {}).get("percent_covered", 0.0) + s_all = ( + max((s.get("totals", {}).get("percent_covered", 0.0) for s in secondaries), default=0.0) + if two + else 0.0 + ) + + out = ["## Primus coverage - %s\n" % title] + if two: + out.append( + "**Unit %.1f%% -> Unit+E2E %.1f%%** (%s statements; excludes %s)\n" + % (_pct(tc, tn), _pct(te, tn), format(tn, ","), excl) + ) + out.append( + "_Including all modules (nothing excluded): Unit %.1f%% -> Unit+E2E %.1f%%._\n" % (p_all, s_all) + ) + out += ["| Module | Stmts | Unit | Unit+E2E |", "|---|--:|--:|--:|"] + 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 |", "|---|--:|--:|--:|"] + + # Top-level groups, sorted by coverage (desc). + for group in sorted(groups, key=lambda 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)) + 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)) + for k, v in sorted(details, key=lambda kv: -_pct(kv[1][0], kv[1][1])): + out.append(row(" `%s`" % k, v[0], v[1], e2e_cov(group, k))) + + out.append(row("TOTAL", tc, tn, te, bold=True)) + return "\n".join(out) + + +def _load(path: str) -> dict: + with open(path) as f: + return json.load(f) + + +def main() -> int: + primary = _load(sys.argv[1]) + title = sys.argv[2] if len(sys.argv) > 2 else "tests" + secondaries = [_load(p) for p in sys.argv[3:]] + print(render(primary, title, secondaries or None)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/ci/junit_summary.py b/tools/ci/junit_summary.py new file mode 100644 index 000000000..b84936674 --- /dev/null +++ b/tools/ci/junit_summary.py @@ -0,0 +1,108 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Render pytest JUnit XML as a Markdown table for the CI run summary. + +GitHub does not render JUnit on its own, so uploading it only yields a download +link. This turns one or more reports into a per-suite pass/fail/error/skip/time +table (plus a collapsible list of failures) to append next to the coverage +table: + + python junit_summary.py --title torch test-reports/*.xml >> "$GITHUB_STEP_SUMMARY" + +Each report is labelled by its filename stem; a missing/unparseable file +renders as a "no report" row instead of failing the step. +""" + +import argparse +import glob +import os +import xml.etree.ElementTree as ET + +COLUMNS = ("tests", "passed", "failed", "errors", "skipped") + + +def parse_file(path): + """Return (label, stats|None, failures) for one JUnit XML file.""" + label = os.path.splitext(os.path.basename(path))[0] + try: + root = ET.parse(path).getroot() + except (ET.ParseError, OSError): + return label, None, [] + + stats = dict.fromkeys(COLUMNS, 0) + stats["time"] = 0.0 + failures = [] + for suite in root.iter("testsuite"): # root is or a lone + stats["tests"] += int(suite.get("tests", 0) or 0) + stats["failed"] += int(suite.get("failures", 0) or 0) + stats["errors"] += int(suite.get("errors", 0) or 0) + stats["skipped"] += int(suite.get("skipped", 0) or 0) + stats["time"] += float(suite.get("time", 0) or 0) + for case in suite.iter("testcase"): + # find(...) "or" is unsafe: an empty Element is falsy, so test explicitly. + bad, kind = case.find("failure"), "failure" + if bad is None: + bad, kind = case.find("error"), "error" + if bad is not None: + name = (case.get("classname", "") + "::" + case.get("name", "")).strip(":") + msg = (bad.get("message") or "").strip().splitlines() + failures.append((name, kind, msg[0] if msg else "")) + stats["passed"] = stats["tests"] - stats["failed"] - stats["errors"] - stats["skipped"] + return label, stats, failures + + +def render(reports, title=None): + lines = ["## Test results - %s\n" % title] if title else [] + lines += [ + "| Suite | Tests | Passed | Failed | Errors | Skipped | Time |", + "|---|--:|--:|--:|--:|--:|--:|", + ] + total = dict.fromkeys(COLUMNS, 0) + total["time"] = 0.0 + failures = [] + for label, st, fails in reports: + if st is 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"]) + ) + 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"]) + ) + lines.append("") + + if failures: + lines.append("
%d failing test(s)\n" % len(failures)) + for label, name, kind, msg in failures: + lines.append(("- `%s` **%s** (%s): %s" % (label, name, kind, msg))[:300]) + lines.append("\n
") + else: + lines.append("_All tests passed._") + return "\n".join(lines) + + +def main(): + ap = argparse.ArgumentParser(description="Render JUnit XML as a Markdown CI summary.") + ap.add_argument("xml", nargs="+", help="JUnit XML file(s) or glob(s).") + ap.add_argument("--title", default=None, help="Optional section title (e.g. torch).") + args = ap.parse_args() + + paths = [] + for pat in args.xml: + paths += sorted(glob.glob(pat)) or [pat] # keep literal -> "no report" row + print(render([parse_file(p) for p in paths], args.title)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/ci/runtime_summary.py b/tools/ci/runtime_summary.py new file mode 100644 index 000000000..dd55ff08c --- /dev/null +++ b/tools/ci/runtime_summary.py @@ -0,0 +1,70 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Render CI stage wall-clock times (a stageseconds TSV) as a Markdown table. + +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. +""" + +import argparse + + +def fmt(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(path): + 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 [] + return rows + + +def render(rows, title=None): + 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)}** |") + return "\n".join(lines) + + +def main(): + ap = argparse.ArgumentParser(description="Render a stageseconds TSV as a Markdown runtime table.") + ap.add_argument("tsv", help="Path to the runtime TSV (stageseconds per line).") + ap.add_argument("--title", default=None, help="Optional section title (e.g. torch).") + 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)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/ci/select_tests.py b/tools/ci/select_tests.py new file mode 100644 index 000000000..cfff17855 --- /dev/null +++ b/tools/ci/select_tests.py @@ -0,0 +1,159 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# 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/ -> tests/unit_tests/), + resolved by walking up to the nearest existing dir; + - E2E suites are auto-discovered from tests/trainer/test__trainer.py; + - a backend is named by its dir (primus/backends/ or examples/). +Fail-safe is the only invariant: anything global, unlocatable, 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 [--e2e] +""" + +import re +import sys +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". +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)/([^/]+)/") + + +def discover_e2e_suites(): + return { + m.group(1) + for p in (ROOT / "tests/trainer").glob("test_*_trainer.py") + if (m := _TRAINER_RE.fullmatch(p.stem)) + } + + +def _is_global(path): + if "/" not in path and path.startswith("requirements") and path.endswith(".txt"): + return True + 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).""" + if _is_global(path): + return ("global", None) + backend = _BACKEND_RE.match(path) # primus/backends// or examples// + if backend: + return ("backend", backend.group(1)) + 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) :])) + 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()] + if not files: + return sorted(suites) + selected = set() + for path in files: + kind, val = classify(path) + if kind == "global": + return sorted(suites) + if kind == "backend": + if val in suites: + selected.add(val) # backend has a trainer -> run its suite + else: + return sorted(suites) # no trainer (bridge/hummingbirdxt/TE) -> all + elif kind == "component": + return sorted(suites) # non-backend source change -> all training + # ignore -> skip + return sorted(selected) + + +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))) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())