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