Update turbo flydsl sparse attn

.github/workflows/ci.yaml

@@ -23,7 +23,7 @@ permissions:
   contents: read
 
 env:
-  PRIMUS_TURBO_COMMIT: 56c789e58f72aaf733b7715f1536be1ed33b69a1 # dev/kyle/flydsl_attn_deepseekv4 (dsv4 sparse-MLA attention + cr=4 fixes + flydsl kernels packaging __init__.py fix)
+  PRIMUS_TURBO_COMMIT: edc8d2ccb0be4888e80ee7c6e765fd3956026a32 # feat: add quantized tensor support for blockwise grouped gemm (#422)
   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)

examples/deepseek-v4/run_deepseek_v4_flash.sh

@@ -12,12 +12,14 @@ if command -v spur >/dev/null 2>&1; then
     export SLURM_ACCOUNT=amd-primus
     # Empty = let the scheduler allocate nodes (skip the hardcoded smci355 default
     # in run_deepseek_v4.sh, whose SLURM_NODELIST uses ${VAR-default} so empty wins).
-    export SLURM_NODELIST=""
+    # Honor an incoming SLURM_NODELIST so callers can pin to specific good nodes
+    # (e.g. to work around a bad node that keeps causing JobLaunchFailure).
+    export SLURM_NODELIST="${SLURM_NODELIST:-}"
     # ABI-4 libionic provider .so to swap into the container at launch (fixes ionic
     # RDMA on AINIC images whose bundled libionic only advertises uverbs ABI 1).
     # The tools/patches/fix_libionic_abi4.sh patch reads this; set empty to disable.
     export PRIMUS_LIBIONIC_SRC_ABI4_SO="bak/ainic/libionic-rdmav34.so.host-abi4/libionic.so.1.0.54.0-149.g3304be71"
-    export NCCL_DEBUG=INFO
+    export NCCL_DEBUG="${NCCL_DEBUG:-}"
     export GLOO_SOCKET_IFNAME=ens3
     export NCCL_SOCKET_IFNAME=ens3
 fi
@@ -32,7 +34,7 @@ export MTP_NUM_LAYERS=${MTP_NUM_LAYERS:-1}
 
 export NNODES=${NNODES:-4}
 
-if [ "$NNODES" -eq 8 ]; then
+if [ "$NNODES" -ge 8 ]; then
     export PRIMUS_TP=${PRIMUS_TP:-1}
     export PRIMUS_PP=${PRIMUS_PP:-8}
     export PRIMUS_EP=${PRIMUS_EP:-8}
@@ -78,7 +80,7 @@ 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}
 export PROFILE=${PROFILE:-False}
-export PRIMUS_EXP_NAME=${PRIMUS_EXP_NAME:-deepseek_v4_flash_proxy_pp${PRIMUS_PP}_ep${PRIMUS_EP}_seq${PRIMUS_SEQ_LENGTH}}
+export PRIMUS_EXP_NAME=${PRIMUS_EXP_NAME:-deepseek_v4_flash_proxy_nodes${NNODES}_pp${PRIMUS_PP}_ep${PRIMUS_EP}_seq${PRIMUS_SEQ_LENGTH}}
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 exec "${SCRIPT_DIR}/run_deepseek_v4.sh" 2>&1 | tee train_flash_ainic.log

primus/backends/megatron/core/transformer/v4_attention_kernels/_turbo_flydsl/__init__.py

@@ -9,10 +9,9 @@
 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.
+native-FlyDSL sparse-MLA attention under ``primus_turbo.flydsl.attention``. This
+is the "turbo API" integration: Primus owns only the thin V4 adapter binding; the
+kernels live in Primus-Turbo.
 
 Public kernel-pair API (identical to the other sparse-MLA backends):
 
@@ -21,21 +20,43 @@
 * ``sparse_mla_bwd_v4_turbo_flydsl(q, kv, o, do, topk, lse, attn_sink=None,
   kv_lora_rank=512, scale=None) -> (dq, dkv, d_sink)``
 
+The upstream module layout / function names changed between Primus-Turbo
+releases, so we bind to whichever is present (the call signatures are identical):
+
+* **new** (>= the ``flydsl_attn`` refactor, e.g. commit ``edc8d2c``):
+  ``primus_turbo.flydsl.attention.sparse_mla_fwd.sparse_mla_fwd_flydsl`` /
+  ``...sparse_mla_bwd.sparse_mla_bwd_flydsl`` (flat modules, no ``_v4`` infix).
+* **old** (<= ``0.3.2.dev*``):
+  ``primus_turbo.flydsl.attention.kernels.sparse_mla_v2.{sparse_mla_fwd_v4_flydsl,
+  sparse_mla_bwd_v4_flydsl}``.
+
 ``primus_turbo`` (with the flydsl sparse-MLA attention) and the ``flydsl`` pip
 package are required; the import fails with a clear message otherwise (handled by
 the lazy loader in :mod:`..` / :func:`load_turbo_attention_backends`).
 """
 
 from __future__ import annotations
 
-from primus_turbo.flydsl.attention.kernels.sparse_mla_v2 import (
-    sparse_mla_bwd_v4_flydsl,
-    sparse_mla_fwd_v4_flydsl,
-)
+try:
+    # New Primus-Turbo layout: flat modules, function names without the _v4 infix.
+    from primus_turbo.flydsl.attention.sparse_mla_bwd import (
+        sparse_mla_bwd_flydsl as _sparse_mla_bwd,
+    )
+    from primus_turbo.flydsl.attention.sparse_mla_fwd import (
+        sparse_mla_fwd_flydsl as _sparse_mla_fwd,
+    )
+except ImportError:
+    # Older Primus-Turbo layout: kernels.sparse_mla_v2 with the _v4 infix.
+    from primus_turbo.flydsl.attention.kernels.sparse_mla_v2 import (
+        sparse_mla_bwd_v4_flydsl as _sparse_mla_bwd,
+    )
+    from primus_turbo.flydsl.attention.kernels.sparse_mla_v2 import (
+        sparse_mla_fwd_v4_flydsl as _sparse_mla_fwd,
+    )
 
 # Turbo-suffixed re-exports so the V4 wrapper / benchmark can bind them without
 # clashing with the in-tree ``_flydsl_v1`` (which has identically-named kernels).
-sparse_mla_fwd_v4_turbo_flydsl = sparse_mla_fwd_v4_flydsl
-sparse_mla_bwd_v4_turbo_flydsl = sparse_mla_bwd_v4_flydsl
+sparse_mla_fwd_v4_turbo_flydsl = _sparse_mla_fwd
+sparse_mla_bwd_v4_turbo_flydsl = _sparse_mla_bwd
 
 __all__ = ["sparse_mla_fwd_v4_turbo_flydsl", "sparse_mla_bwd_v4_turbo_flydsl"]

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

@@ -43,10 +43,15 @@
 
 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)",
-)
+# The upstream module layout changed across releases; accept either (new flat
+# ``sparse_mla_fwd`` module, or the old ``kernels.sparse_mla_v2`` package).
+try:
+    import primus_turbo.flydsl.attention.sparse_mla_fwd  # noqa: F401
+except ImportError:
+    pytest.importorskip(
+        "primus_turbo.flydsl.attention.kernels.sparse_mla_v2",
+        reason="installed primus_turbo has no flydsl sparse-MLA attention (turbo backend)",
+    )
 
 _ARCH = torch.cuda.get_device_properties(0).gcnArchName
 if "gfx950" not in _ARCH: