ci: Fix build-and-test-wheels jobs for arm

.github/workflows/build-test-publish-wheel.yml

@@ -66,14 +66,14 @@ jobs:
           GITHUB_RUN_ID: ${{ github.run_id }}
           SKIPPING_IS_ALLOWED: false
         run: |
-          FAILED_JOBS=$(gh run view $GITHUB_RUN_ID --json jobs --jq '[.jobs[] | select(.status == "completed" and .conclusion != "success")] | length') || echo 0
+          FAILED_JOBS=$(gh run view $GITHUB_RUN_ID --json jobs --jq '[.jobs[] | select(.status == "completed" and .conclusion != "success" and (.name | test("build-and-test-wheels")))] | length') || echo 0
 
           if [ "${FAILED_JOBS:-0}" -eq 0 ] || [ "$SKIPPING_IS_ALLOWED" == "true" ]; then
-              echo "✅ All previous jobs completed successfully"
+              echo "✅ All build-and-test-wheels jobs completed successfully"
               exit 0
           else
-              echo "❌ Found $FAILED_JOBS failed job(s)"
+              echo "❌ Found $FAILED_JOBS failed build-and-test-wheels job(s)"
               # Show which jobs failed
-              gh run view $GITHUB_RUN_ID --json jobs --jq '.jobs[] | select(.status == "completed" and .conclusion != "success") | .name'
+              gh run view $GITHUB_RUN_ID --json jobs --jq '.jobs[] | select(.status == "completed" and .conclusion != "success" and (.name | test("build-and-test-wheels"))) | .name'
               exit 1
           fi

megatron/core/inference/contexts/attention_context/triton/tensor_ops.py

@@ -3,8 +3,21 @@
 from typing import Optional
 
 import torch
-import triton  # type: ignore
-import triton.language as tl  # type: ignore
+
+try:
+    import triton
+    import triton.language as tl
+
+    HAVE_TRITON = True
+except ImportError:
+    from unittest.mock import MagicMock
+
+    from megatron.core.utils import null_decorator
+
+    triton = MagicMock()
+    triton.jit = null_decorator
+    tl = MagicMock()
+    HAVE_TRITON = False
 
 
 @triton.jit

megatron/core/inference/contexts/fused_kv_append_kernel.py

@@ -1,9 +1,22 @@
 # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 
-import triton
-import triton.language as tl
 from torch import Tensor
 
+try:
+    import triton
+    import triton.language as tl
+
+    HAVE_TRITON = True
+except ImportError:
+    from unittest.mock import MagicMock
+
+    from megatron.core.utils import null_decorator
+
+    triton = MagicMock()
+    triton.jit = null_decorator
+    tl = MagicMock()
+    HAVE_TRITON = False
+
 
 @triton.jit
 def _append_kv_cache_kernel(

megatron/core/inference/quantization/mxfp8_quantize.py

@@ -13,8 +13,21 @@
 """
 
 import torch
-import triton
-import triton.language as tl
+
+try:
+    import triton
+    import triton.language as tl
+
+    HAVE_TRITON = True
+except ImportError:
+    from unittest.mock import MagicMock
+
+    from megatron.core.utils import null_decorator
+
+    triton = MagicMock()
+    triton.jit = null_decorator
+    tl = MagicMock()
+    HAVE_TRITON = False
 
 
 def _ceil_div(a, b):

megatron/core/ssm/ops/causal_conv1d_triton.py

@@ -5,8 +5,21 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
-import triton
-import triton.language as tl
+
+try:
+    import triton
+    import triton.language as tl
+
+    HAVE_TRITON = True
+except ImportError:
+    from unittest.mock import MagicMock
+
+    from megatron.core.utils import null_decorator
+
+    triton = MagicMock()
+    triton.jit = null_decorator
+    tl = MagicMock()
+    HAVE_TRITON = False
 
 
 @triton.jit

megatron/core/ssm/ops/mamba_ssm.py

@@ -6,12 +6,25 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
-import triton
-import triton.language as tl
 from packaging import version
 
-TRITON3 = version.parse(triton.__version__) >= version.parse("3.0.0")
+try:
+    import triton
+    import triton.language as tl
 
+    TRITON3 = version.parse(triton.__version__) >= version.parse("3.0.0")
+
+    HAVE_TRITON = True
+except ImportError:
+    from unittest.mock import MagicMock
+
+    from megatron.core.utils import null_decorator
+
+    triton = MagicMock()
+    triton.jit = null_decorator
+    tl = MagicMock()
+    HAVE_TRITON = False
+    TRITON3 = False
 
 if TRITON3:
 
@@ -20,7 +33,7 @@ def softplus(dt):
         """Optimized softplus."""
         return tl.math.log(tl.math.exp(dt) + 1)
 
-else:
+elif HAVE_TRITON:
 
     @triton.jit
     def softplus(dt):

megatron/core/transformer/custom_layers/batch_invariant_kernels.py

@@ -12,8 +12,21 @@
 from typing import Any, Dict, List, Optional
 
 import torch
-import triton
-import triton.language as tl
+
+try:
+    import triton
+    import triton.language as tl
+
+    HAVE_TRITON = True
+except ImportError:
+    from unittest.mock import MagicMock
+
+    from megatron.core.utils import null_decorator
+
+    triton = MagicMock()
+    triton.jit = null_decorator
+    tl = MagicMock()
+    HAVE_TRITON = False
 
 __all__ = [
     "set_batch_invariant_mode",