Update turbo flydsl sparse attn

[wenxie-amd]

No description provided.

[Copilot]

Pull request overview

Updates the DeepSeek-V4 Flash run wrapper defaults and bumps the Primus-Turbo commit used by CI, with the intent of picking up newer flydsl sparse attention support.

Changes:

• Adjust DeepSeek-V4 Flash SLURM env defaults (honor incoming node list; tweak NCCL_DEBUG; broaden NNODES condition).
• Update experiment naming to include node count for easier run identification.
• Bump PRIMUS_TURBO_COMMIT in CI, but also comment out the torch/jax unit-test + coverage jobs.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
examples/deepseek-v4/run_deepseek_v4_flash.sh	Updates SLURM/NCCL defaults and run naming for DeepSeek-V4 Flash launcher wrapper.
.github/workflows/ci.yaml	Updates the Primus-Turbo pin, but disables unit-test/coverage jobs by commenting them out.

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (1)

.github/workflows/ci.yaml:212

• This change comments out the entire run-unittest-torch, run-unittest-jax, and coverage-summary jobs, effectively disabling unit/E2E testing and coverage reporting for the workflow. That is a major behavior change and doesn’t align with the PR’s stated scope (updating the turbo FlyDSL sparse attention). Please restore these jobs (or move them to a replacement workflow) so CI continues to execute tests/coverage.

#  run-unittest-torch:
#    permissions:
#      contents: read
#      actions: read # let "Write runtime summary" list this job's own steps for auto-discovered timing
#    env:
#      PRIMUS_WORKDIR: /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 }}

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (3)

tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_flydsl_attention.py:54

• The skip-logic only probes the new forward module (and falls back on any ImportError), but the turbo adapter requires both fwd and bwd. This can make the test fail later with a confusing ImportError if only one of the flat modules is present, and it can also silently fall back to the old layout when the new module exists but fails to import due to a missing internal dependency.

# The upstream module layout changed across releases; accept either (new flat
# ``sparse_mla_fwd`` module, or the old ``kernels.sparse_mla_v2`` package).
try:
    import primus_turbo.flydsl.attention.sparse_mla_fwd  # noqa: F401
except ImportError:
    pytest.importorskip(
        "primus_turbo.flydsl.attention.kernels.sparse_mla_v2",
        reason="installed primus_turbo has no flydsl sparse-MLA attention (turbo backend)",
    )

primus/backends/megatron/core/transformer/v4_attention_kernels/_turbo_flydsl/init.py:55

• The broad except ImportError will also catch import failures originating from missing internal dependencies of the new Primus-Turbo modules (e.g., a missing shared library), and then silently fall back to the old layout. That can hide a broken installation and make failures harder to diagnose. Prefer falling back only when the new modules or symbols are absent, and re-raise unexpected import errors.

try:
    # New Primus-Turbo layout: flat modules, function names without the _v4 infix.
    from primus_turbo.flydsl.attention.sparse_mla_bwd import (
        sparse_mla_bwd_flydsl as _sparse_mla_bwd,
    )
    from primus_turbo.flydsl.attention.sparse_mla_fwd import (
        sparse_mla_fwd_flydsl as _sparse_mla_fwd,
    )
except ImportError:
    # Older Primus-Turbo layout: kernels.sparse_mla_v2 with the _v4 infix.
    from primus_turbo.flydsl.attention.kernels.sparse_mla_v2 import (
        sparse_mla_bwd_v4_flydsl as _sparse_mla_bwd,
    )
    from primus_turbo.flydsl.attention.kernels.sparse_mla_v2 import (
        sparse_mla_fwd_v4_flydsl as _sparse_mla_fwd,
    )

examples/deepseek-v4/run_deepseek_v4_flash.sh:41

• With set -u, if NNODES is set to a value other than 4 or >=8, this script never exports PRIMUS_PP / PRIMUS_EP, but later interpolates them into PRIMUS_EXP_NAME, causing an unbound-variable crash. Add an explicit fallback path that either sets defaults or requires callers to provide the missing env vars.

if [ "$NNODES" -ge 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

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (2)

tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_flydsl_attention.py:52

• The skip guard only checks that the new forward module exists. If a Primus-Turbo install has sparse_mla_fwd but is missing the matching backward module, the test module will import successfully here but then fail later when v4_attention_turbo tries to bind the backward kernel. Import both modules (or both symbols) in the try-block so the test cleanly skips when either half is unavailable.

try:
    import primus_turbo.flydsl.attention.sparse_mla_fwd  # noqa: F401
except ImportError:
    pytest.importorskip(
        "primus_turbo.flydsl.attention.kernels.sparse_mla_v2",

primus/backends/megatron/core/transformer/v4_attention_kernels/_turbo_flydsl/init.py:44

• Catching a broad ImportError here will also swallow ImportErrors raised from inside the new Primus-Turbo modules (e.g., failing to load a required shared library), and then silently try the old layout. That makes real import/load failures harder to debug and can change behavior. Prefer only falling back when the new modules/symbols are actually absent (ModuleNotFoundError / AttributeError).

try:
    # New Primus-Turbo layout: flat modules, function names without the _v4 infix.
    from primus_turbo.flydsl.attention.sparse_mla_bwd import (
        sparse_mla_bwd_flydsl as _sparse_mla_bwd,
    )

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (3)

tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_flydsl_attention.py:52

• The except ImportError here will also catch import failures inside the new turbo modules (e.g., missing shared libraries or a transitive dependency) and will then importorskip(...), causing the test to be skipped instead of failing. That can hide real regressions in primus_turbo/flydsl packaging or runtime dependencies. Consider only falling back when the module itself is missing, and also validate both fwd and bwd modules for the new layout.

try:
    import primus_turbo.flydsl.attention.sparse_mla_fwd  # noqa: F401
except ImportError:
    pytest.importorskip(
        "primus_turbo.flydsl.attention.kernels.sparse_mla_v2",

primus/backends/megatron/core/transformer/v4_attention_kernels/_turbo_flydsl/init.py:44

• Catching a broad ImportError here can mask real failures inside the new Primus-Turbo modules (e.g., missing runtime deps) and silently fall back to the old layout, which makes failures harder to debug and can lead to unintentionally running the wrong kernel implementation. It’s safer to fall back only when the new modules/functions are actually missing, and re-raise other ImportErrors.

try:
    # New Primus-Turbo layout: flat modules, function names without the _v4 infix.
    from primus_turbo.flydsl.attention.sparse_mla_bwd import (
        sparse_mla_bwd_flydsl as _sparse_mla_bwd,
    )

examples/deepseek-v4/run_deepseek_v4_flash.sh:22

• Exporting NCCL_DEBUG as an empty string changes semantics vs leaving it unset (some NCCL builds treat an empty-but-present value as invalid and may warn or behave unexpectedly). If the intent is “don’t override unless the caller set it”, only export when it’s non-empty/present.

    export NCCL_DEBUG="${NCCL_DEBUG:-}"

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (2)

tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_flydsl_attention.py:52

• The import check only verifies sparse_mla_fwd and catches a broad ImportError, which can (a) let the test run even if sparse_mla_bwd is missing and (b) accidentally skip when the new module exists but fails due to an internal import error. Consider importing both fwd+bwd and only falling back when the module itself is missing; otherwise re-raise to surface real breakages.

try:
    import primus_turbo.flydsl.attention.sparse_mla_fwd  # noqa: F401
except ImportError:
    pytest.importorskip(
        "primus_turbo.flydsl.attention.kernels.sparse_mla_v2",

primus/backends/megatron/core/transformer/v4_attention_kernels/_turbo_flydsl/init.py:44

• The layout-selection try/except ImportError will also fall back on internal import failures (e.g., a missing dependency inside the new modules), which can mask real breakages by silently binding the old path. Prefer only falling back when the new modules themselves are missing; for other import errors, re-raise so load_turbo_attention_backends() reports the true root cause.

try:
    # New Primus-Turbo layout: flat modules, function names without the _v4 infix.
    from primus_turbo.flydsl.attention.sparse_mla_bwd import (
        sparse_mla_bwd_flydsl as _sparse_mla_bwd,
    )

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.

Comments suppressed due to low confidence (1)

tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_flydsl_attention.py:52

• The test only checks for primus_turbo.flydsl.attention.sparse_mla_fwd in the new layout, but the Primus adapter imports both sparse_mla_fwd and sparse_mla_bwd. If a Primus-Turbo build is partially installed (fwd present, bwd missing), this test will not skip and will later fail on import. Import/skip should validate both modules for the new layout.

try:
    import primus_turbo.flydsl.attention.sparse_mla_fwd  # noqa: F401
except ImportError:
    pytest.importorskip(
        "primus_turbo.flydsl.attention.kernels.sparse_mla_v2",

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (3)

tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_flydsl_attention.py:54

• The fallback check currently catches any ImportError from importing primus_turbo.flydsl.attention.sparse_mla_fwd. That can incorrectly skip this test when the new module exists but fails to import due to a missing dependency or an internal import error (in that case, the test should fail rather than silently falling back / skipping). Limit the fallback to the case where the module itself is missing (ModuleNotFoundError for that module).

try:
    import primus_turbo.flydsl.attention.sparse_mla_fwd  # noqa: F401
except ImportError:
    pytest.importorskip(
        "primus_turbo.flydsl.attention.kernels.sparse_mla_v2",
        reason="installed primus_turbo has no flydsl sparse-MLA attention (turbo backend)",
    )

primus/backends/megatron/core/transformer/v4_attention_kernels/_turbo_flydsl/init.py:55

• This import-compat shim falls back on any ImportError from the new Primus-Turbo modules. That can mask real errors (e.g., a missing dependency imported inside sparse_mla_fwd/sparse_mla_bwd) and silently route to the old kernel implementation. Prefer falling back only when the new modules are absent or when the expected symbols are missing; otherwise re-raise to surface the actual problem.

try:
    # New Primus-Turbo layout: flat modules, function names without the _v4 infix.
    from primus_turbo.flydsl.attention.sparse_mla_bwd import (
        sparse_mla_bwd_flydsl as _sparse_mla_bwd,
    )
    from primus_turbo.flydsl.attention.sparse_mla_fwd import (
        sparse_mla_fwd_flydsl as _sparse_mla_fwd,
    )
except ImportError:
    # Older Primus-Turbo layout: kernels.sparse_mla_v2 with the _v4 infix.
    from primus_turbo.flydsl.attention.kernels.sparse_mla_v2 import (
        sparse_mla_bwd_v4_flydsl as _sparse_mla_bwd,
    )
    from primus_turbo.flydsl.attention.kernels.sparse_mla_v2 import (
        sparse_mla_fwd_v4_flydsl as _sparse_mla_fwd,
    )

.github/workflows/ci.yaml:294

• This line's nested quoting looks malformed (print('\''triton'\'', ...)) and is very likely to break the bash -lc command (or produce an invalid Python snippet). Use consistent quoting (e.g., escape double-quotes inside the Python one-liner) to keep the step reliable.

          docker exec "${UT_CONTAINER}" bash -lc 'python3 --version; python3 -c "import triton; print('\''triton'\'', triton.__version__)" 2>/dev/null || true'

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (2)

tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_flydsl_attention.py:54

• The broad except ImportError will also catch ImportErrors raised inside primus_turbo.flydsl.attention.sparse_mla_fwd (e.g., missing symbols/deps) and then silently fall back to importorskip, which can incorrectly skip the test and hide a real regression. Only fall back when the module itself (or primus_turbo) is missing; otherwise let the import error surface.

try:
    import primus_turbo.flydsl.attention.sparse_mla_fwd  # noqa: F401
except ImportError:
    pytest.importorskip(
        "primus_turbo.flydsl.attention.kernels.sparse_mla_v2",
        reason="installed primus_turbo has no flydsl sparse-MLA attention (turbo backend)",
    )

primus/backends/megatron/core/transformer/v4_attention_kernels/_turbo_flydsl/init.py:55

• Catching a broad ImportError here can mask real import-time failures in the new Primus-Turbo modules (e.g., missing symbols inside sparse_mla_fwd) and then silently fall back to the old layout. Only fall back when the new modules (or primus_turbo) are actually missing; otherwise let the error surface so failures are diagnosable.

try:
    # New Primus-Turbo layout: flat modules, function names without the _v4 infix.
    from primus_turbo.flydsl.attention.sparse_mla_bwd import (
        sparse_mla_bwd_flydsl as _sparse_mla_bwd,
    )
    from primus_turbo.flydsl.attention.sparse_mla_fwd import (
        sparse_mla_fwd_flydsl as _sparse_mla_fwd,
    )
except ImportError:
    # Older Primus-Turbo layout: kernels.sparse_mla_v2 with the _v4 infix.
    from primus_turbo.flydsl.attention.kernels.sparse_mla_v2 import (
        sparse_mla_bwd_v4_flydsl as _sparse_mla_bwd,
    )
    from primus_turbo.flydsl.attention.kernels.sparse_mla_v2 import (
        sparse_mla_fwd_v4_flydsl as _sparse_mla_fwd,
    )

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (3)

primus/backends/megatron/core/transformer/v4_attention_kernels/_turbo_flydsl/init.py:48

• Catching a broad ImportError here can mask real import failures inside the new Primus-Turbo modules (e.g., missing dependency or an exception raised during module import) and silently fall back to the old path. If the intent is only to fall back when the module is absent, catch ModuleNotFoundError instead so other ImportErrors surface clearly.

except ImportError:

tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_flydsl_attention.py:50

• This except ImportError will also catch ImportErrors thrown inside primus_turbo.flydsl.attention.sparse_mla_fwd during import, which would incorrectly downgrade to the old layout and potentially skip a real regression. If the goal is compatibility based on module presence, catch ModuleNotFoundError so genuine import failures are not hidden.

except ImportError:

.github/workflows/ci.yaml:541

• tools/ci/runtime_summary.py relies on Actions-provided env vars (GITHUB_REPOSITORY, GITHUB_RUN_ID, GITHUB_JOB, RUNNER_NAME). Running it inside the container without passing these makes it return early and produce an empty runtime summary.

          docker exec -e GITHUB_WORKSPACE="${CWS}" -e GITHUB_TOKEN -e RUNNER_TEMP -w "${CWS}" "${UT_CONTAINER}" \
            bash -lc 'python tools/ci/runtime_summary.py --title torch' >> "$GITHUB_STEP_SUMMARY" || true

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (2)

tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_flydsl_attention.py:50

• except ImportError here will also catch ImportErrors raised inside primus_turbo.flydsl.attention.sparse_mla_fwd (e.g., a missing dependency). That can silently fall back to the old module path and mask a real installation/runtime failure, making test results harder to interpret.

except ImportError:

primus/backends/megatron/core/transformer/v4_attention_kernels/_turbo_flydsl/init.py:48

• The broad except ImportError fallback can mask real import-time failures in the new Primus-Turbo modules (e.g., a missing runtime dependency inside sparse_mla_fwd/bwd). If the module exists but is broken, we should surface that error instead of silently trying the legacy path.

except ImportError:

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (2)

primus/backends/megatron/core/transformer/v4_attention_kernels/_turbo_flydsl/init.py:44

• The compatibility import uses except ImportError, which will also catch ImportErrors raised inside the new Primus-Turbo modules (e.g., missing runtime deps / shared libs) and silently fall back to the legacy path, masking the real failure and potentially producing a misleading error. Limit the fallback to ModuleNotFoundError for the specific new-module paths and re-raise otherwise (similar to the repo’s TE compatibility imports).

try:
    # New Primus-Turbo layout: flat modules, function names without the _v4 infix.
    from primus_turbo.flydsl.attention.sparse_mla_bwd import (
        sparse_mla_bwd_flydsl as _sparse_mla_bwd,
    )

tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_flydsl_attention.py:52

• This version-compat import guard catches all ImportError, which can incorrectly skip the test when the new module exists but fails to import due to an internal ImportError (e.g., missing runtime dependency). Narrow the fallback to ModuleNotFoundError for the specific missing-module cases so real import failures surface as test failures.

try:
    import primus_turbo.flydsl.attention.sparse_mla_fwd  # noqa: F401
except ImportError:
    pytest.importorskip(
        "primus_turbo.flydsl.attention.kernels.sparse_mla_v2",