fix(packaging): make triton optional at import

tests/unit_tests/test_minimal_async_ep_kernels.py

@@ -6,8 +6,10 @@
 
 import unittest
 
+import pytest
 import torch
-import triton.language as tl
+
+tl = pytest.importorskip("triton.language")
 
 from torchtitan.distributed.minimal_async_ep.kernels import (
     _copy_rows_to_peer_ptrs_kernel,

tests/unit_tests/test_tp_kv_heads_validation.py

@@ -16,14 +16,6 @@
 from types import SimpleNamespace
 
 try:
-    import sys
-    from unittest.mock import MagicMock
-
-    # triton is Linux-only; stub it so MoE kernel imports don't crash on Windows
-    if "triton" not in sys.modules:
-        sys.modules["triton"] = MagicMock()
-        sys.modules["triton.language"] = MagicMock()
-
     from torchtitan.config import ParallelismConfig
     from torchtitan.models.common import (
         ComplexRoPE,

torchtitan/distributed/minimal_async_ep/api.py

@@ -28,16 +28,21 @@
 import torch.distributed as dist
 import torch.distributed._symmetric_memory as symm_mem
 
-from torchtitan.distributed.minimal_async_ep.kernels import (
-    copy_full_counts_to_peers_kernel,
-    copy_rows_to_peers_kernel,
-    expand_topk_grad_kernel,
-    fill_combine_metadata_kernel,
-    fill_dispatch_metadata_kernel,
-    invert_flat_indices_kernel,
-    reduce_topk_slots_kernel,
-    topk_scores_grad_kernel,
-)
+try:
+    from torchtitan.distributed.minimal_async_ep.kernels import (
+        copy_full_counts_to_peers_kernel,
+        copy_rows_to_peers_kernel,
+        expand_topk_grad_kernel,
+        fill_combine_metadata_kernel,
+        fill_dispatch_metadata_kernel,
+        invert_flat_indices_kernel,
+        reduce_topk_slots_kernel,
+        topk_scores_grad_kernel,
+    )
+except ImportError:
+    # No triton: the shared model code imports this module, but the CUDA-only
+    # kernels are never launched.
+    pass
 from torchtitan.tools.logging import logger
 
 

torchtitan/overrides/fused_swiglu.py

@@ -53,9 +53,10 @@
 
 import spmd_types as spmd
 import torch
-import triton
-import triton.language as tl
 
+# Without triton, triton_compat's jit stubs raise only when a kernel is launched,
+# so this module still imports.
+from torch._inductor.runtime.triton_compat import tl, triton
 from torch.distributed.tensor import DTensor
 from torch.distributed.tensor.experimental import local_map