AI計算を速める新技術を解説

このチュートリアルでは、さまざまなハードウェア条件で動作する実用的なColabワークフローを構築することで、TileGym GPUプログラミングを探求します。我々はまず、利用可能なCUDA環境を調査し、NVIDIA cuTileが直接実行されるかを確認し、標準のColab GPUに必要なcuTileスタックがない場合はTritonにフォールバックします。この設定を通じて、タイルプログラミングの核となるアイデアを学びます。それは、一度に1つのスレッドのコードを書く代わりに、データタイル全体を操作し、それらをカーネルにロードし、効率的に計算し、結果を保存し直すというものです。このモデルを使用して、ベクトル加算、fused GELU、行ごとのsoftmax、タイル化された行列乗算、およびflash attentionを実装し、それぞれの結果をPyTorchと比較して正確性とベンチマークを行います。CUDA環境の調査

コードをコピー

コピー済み

別のブラウザを使用

import os, sys, math, time, textwrap

def rule(t=""):

print("\n" + "=" * 78)

if t: print(t)

print("=" * 78)

rule("0. ENVIRONMENT PROBE")

try:

import torch

except ImportError:

print("Installing torch ...")

os.system(f"{sys.executable} -m pip install -q torch")

import torch

HAS_CUDA = torch.cuda.is_available()

DEV = "cuda" if HAS_CUDA else "cpu"

cc = (0, 0)

if HAS_CUDA:

cc = torch.cuda.get_device_capability()

print(f"GPU : {torch.cuda.get_device_name(0)}")

print(f"Compute capability : {cc[0]}.{cc[1]}")

print(f"Torch CUDA runtime : {torch.version.cuda}")

print(f"Driver / torch : {torch.__version__}")

else:

print("No CUDA GPU found. In Colab: Runtime -> Change runtime type -> GPU (T4).")

print("The tutorial will still run its correctness math on CPU where possible.")

CUTILE_HW_OK = HAS_CUDA and cc[0] >= 8

CUDA_MAJOR = int((torch.version.cuda or "0").split(".")[0]) if HAS_CUDA else 0

CUTILE_TOOLKIT_OK = CUDA_MAJOR >= 13

rule("1. ATTEMPTING REAL cuTile (NVIDIA CUDA Tile) BACKEND")

ct = None

CUTILE_READY = False

if CUTILE_HW_OK and CUTILE_TOOLKIT_OK:

try:

import cuda.tile as ct

CUTILE_READY = True

print("cuda.tile is already importable.")

except Exception:

print("Instal