From dbb71e45f3d3689d4a91a74d4088b87b97fc79c1 Mon Sep 17 00:00:00 2001
From: ahoffman-aws <183871518+ahoffman-aws@users.noreply.github.com>
Date: Wed, 11 Nov 2026 01:21:23 +0000
Subject: [PATCH] fix(packaging): make triton optional at import

---
 .../test_minimal_async_ep_kernels.py          |  4 ++-
 .../unit_tests/test_tp_kv_heads_validation.py |  8 ------
 .../distributed/minimal_async_ep/api.py       | 25 +++++++++++--------
 torchtitan/overrides/fused_swiglu.py          |  5 ++--
 4 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/tests/unit_tests/test_minimal_async_ep_kernels.py b/tests/unit_tests/test_minimal_async_ep_kernels.py
index cbb0572fd87..6672ebcc181 100644
--- a/tests/unit_tests/test_minimal_async_ep_kernels.py
+++ b/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,
diff --git a/tests/unit_tests/test_tp_kv_heads_validation.py b/tests/unit_tests/test_tp_kv_heads_validation.py
index 0046374b7fc..5c385359fdc 100644
--- a/tests/unit_tests/test_tp_kv_heads_validation.py
+++ b/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,
diff --git a/torchtitan/distributed/minimal_async_ep/api.py b/torchtitan/distributed/minimal_async_ep/api.py
index b8b2affea0f..1d1bc134dd8 100644
--- a/torchtitan/distributed/minimal_async_ep/api.py
+++ b/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
 
 
diff --git a/torchtitan/overrides/fused_swiglu.py b/torchtitan/overrides/fused_swiglu.py
index bdfa8ac22f0..7f3a9132066 100644
--- a/torchtitan/overrides/fused_swiglu.py
+++ b/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