WindowsでTritonをインストールしたにもかかわらず、torch_has_triton()Falseを返してしまう問題に遭遇することがあります。この記事では、その原因と解決方法を解説します。


Tritonとは?

 GPUカーネルをPythonで書けるようにするコンパイラ/言語です。
通常、GPUで高速な処理をするにはCUDA(C++)でカーネルを書く必要がありますが、TritonはそれをPythonに近い書き方で実現できます。

具体的にやっていること
行列演算・Attention・活性化関数などのGPUカーネルを自動生成
メモリアクセスの最適化(メモリ帯域を無駄なく使う)
FlashAttentionなどの高効率な実装もTritonで書かれている
まとめると
「PythonでCUDAカーネルを書けるようにして、torch.compile() の高速化を支えるもの」

Tritonが使えないと torch.compile() が効かず、モデルの推論・学習が遅くなります。

発生する症状

以下のようなコードで triton_is_availableFalse になってしまいます。

try:
    if torch._dynamo.config.disable:
        triton_is_available = False
    else:
        from torch.utils._triton import has_triton as torch_has_triton
        triton_is_available = torch_has_triton()
except Exception:
    triton_is_available = False

has_triton() の内部処理

SDNQライブラリのhas_triton() は単純にtritonがインポートできるかだけを見ているわけではありません。ソースコードを確認すると、以下のような処理をしています。

@functools.lru_cache(None)
def has_triton() -> bool:
    if not has_triton_package():
        return False
    from torch._dynamo.device_interface import get_interface_for_device

    def cuda_extra_check(device_interface):
        return device_interface.Worker.get_device_properties().major >= 7

    def cpu_extra_check(device_interface):
        import triton.backends
        return "cpu" in triton.backends.backends

    def _return_true(device_interface):
        return True

    triton_supported_devices = {
        "cuda": cuda_extra_check,
        "xpu": _return_true,
        "cpu": cpu_extra_check,
    }

    def is_device_compatible_with_triton():
        for device, extra_check in triton_supported_devices.items():
            device_interface = get_interface_for_device(device)
            if device_interface.is_available() and extra_check(device_interface):
                return True
        return False

    return is_device_compatible_with_triton()

つまり、以下の条件をすべて満たす必要があります。

  • tritonパッケージがインストールされている
  • CUDAが利用可能、かつGPUのCompute Capability(major版)が 7以上 である

原因の切り分け方

以下のコードを実行して、どこで失敗しているか確認しましょう。

import torch
from torch._dynamo.device_interface import get_interface_for_device

# CUDAが使えるか確認
cuda_interface = get_interface_for_device("cuda")
print("CUDA available:", cuda_interface.is_available())

# GPUのCompute Capabilityを確認(7以上が必要)
if cuda_interface.is_available():
    props = cuda_interface.Worker.get_device_properties()
    print("GPU major version:", props.major)

# tritonのバックエンド確認
import triton.backends
print("triton backends:", triton.backends.backends)

GPU世代ごとの対応状況

major値 GPU世代 代表製品 Triton対応
6 Pascal GTX 10xxシリーズ ❌ 非対応
7 Volta / Turing RTX 20xxシリーズ ✅ 対応
8 Ampere RTX 30xxシリーズ ✅ 対応
9 Ada Lovelace RTX 40xxシリーズ ✅ 対応

解決方法:torchバージョンに合わせたtriton-windowsをインストールする

最も多い原因は、インストールしている triton-windows のバージョンがPyTorchのバージョンと一致していないことです。

手順1:既存のtritonを削除する

pip uninstall triton triton-windows -y

手順2:PyTorchのバージョンを確認する

python -c "import torch; print(torch.__version__)"

手順3:torchバージョンに対応したtriton-windowsをインストールする

PyTorchのバージョンに合わせて、対応するtriton-windowsをインストールします。バージョンは  triton-windows のページ で確認できます。

# 例:torch 2.3.x の場合
pip install triton-windows==2.3.*

# 例:torch 2.4.x の場合
pip install triton-windows==2.4.*

手順4:動作確認

import torch
from torch.utils._triton import has_triton
print(has_triton())  # True になればOK

補足:PyTorch公式のwhlを使う方法

PyTorch 2.x以降では、PyTorch公式サイト からCUDAバージョンに合ったwhlをインストールすることで、tritonが自動的にバンドルされるため、別途インストールが不要になります。(ただしtriton-windowsではなくLinux版のtritonとなります)

# CUDA 12.1の場合の例
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

まとめ

  • torch_has_triton() は単純なインポート確認ではなく、CUDAの利用可否GPUのCompute Capability(major >= 7) もチェックしている
  • Windowsで triton-windows を使う場合は、PyTorchのバージョンと対応するバージョンをインストールする必要がある
  • 可能であれば、PyTorch公式のwhlを使うとtritonが自動でバンドルされるため管理が楽になる