diff --git a/.github/workflows/build_kernel.yaml b/.github/workflows/build_kernel.yaml index 1ed9ac6a..953cef32 100644 --- a/.github/workflows/build_kernel.yaml +++ b/.github/workflows/build_kernel.yaml @@ -64,6 +64,7 @@ jobs: relu-kernel-cpu relu-backprop-compile-kernel relu-triton-kernel + gemm-triton-autotune-kernel silu-and-mul-kernel test: diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 1b57b8da..c55f0c2e 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -68,6 +68,10 @@ - local: cli-skills title: Install agent skills title: kernel-builder CLI +- sections: + - local: builder/triton-autotune + title: Ship Triton autotune configurations + title: Advanced practices - sections: - local: kernel-requirements title: Kernel requirements diff --git a/docs/source/builder/triton-autotune.md b/docs/source/builder/triton-autotune.md new file mode 100644 index 00000000..04f61998 --- /dev/null +++ b/docs/source/builder/triton-autotune.md @@ -0,0 +1,177 @@ +# Ship Triton autotune configurations + +Triton kernels typically have parameters, such as tile sizes, number of warps and +number of pipeline stages, whose optimal values depend on the GPU and the +problem shape. Autotuning finds good values for these parameters, but doing +it at runtime (e.g. with the +[`@triton.autotune`](https://triton-lang.org/main/python-api/generated/triton.autotune.html) +decorator) re-benchmarks every candidate configuration in each new process. + + +However, one can run the tuner for the GPU models they like and store the +best found configurations as files for using them later. This effectively +reduces the potentially costly tuning time. + +`kernel-builder` support this by packaging these configurations as JSON files +with the kernel. At runtime, the kernel looks up the configuration for the +current GPU and shape and falls back to sensible defaults when there is no +matching configuration. + +This is how, for example, the vLLM fused MoE kernel ships on the Hub: the +[`RedHatAI/moe`](https://huggingface.co/RedHatAI/moe/tree/main/torch-ext/moe/configs) +repository contains a `configs` directory with tuned configurations for many +GPUs. This page walks through a small, complete example of the same pattern: +the [`gemm-triton-autotune`](https://github.com/huggingface/kernels/tree/main/examples/kernels/gemm-triton-autotune) +example kernel, a Triton GEMM published as +[`kernels-test/gemm-triton-autotune`](https://huggingface.co/kernels-test/gemm-triton-autotune). + +## Shipping data files with a kernel + +Since autotune files are plain JSON files, we can store them anywhere inside the +kernels main Python sources in `torch-ext/`. For this example, +we will use `torch-ext//configs/`. By default, only `py` and `pyi` +files are picked up from the kenel's Python source directory, so add `json` to the +[`pyext` option](writing-kernels.md#torch-noarch) in `build.toml`: + +```toml +[general] +name = "gemm-triton-autotune" +version = 1 +edition = 5 +license = "Apache-2.0" +backends = ["cuda", "rocm", "xpu"] + +[general.hub] +repo-id = "kernels-test/gemm-triton-autotune" + +[torch-noarch] +pyext = ["json", "py"] +``` + +## Configuration file layout + +A GEMM computes `(M, K) @ (K, N)`. For a model, the weight dimensions `N` +and `K` are known ahead of time, while `M` (e.g. the number of tokens) +varies at runtime. The example therefore stores one file per `(N, K)` shape +and GPU, following the same naming convention as the MoE kernel: + +``` +configs/N=4096,K=4096,device_name=NVIDIA_L4.json +configs/N=14336,K=4096,device_name=NVIDIA_L4.json +``` + +Each file maps an `M` value to the best configuration that the autotuner +found for that `M`: + +```json +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 3 + } +} +``` + +Since the device name is part of the file name, configurations tuned for one +GPU are never applied to another. A configuration that would exceed the +resources of a smaller GPU (e.g. shared memory) is therefore harmless to +ship. + +## Looking up configurations at runtime + +At kernel launch, the kernel checks whether a configuration file exists for +the current device and shape. If it does, the configuration with the nearest +tuned `M` is used; otherwise the kernel falls back to a conservative default +and logs a warning. The lookup is cached, so the file is read at most once +per process: + +```python +@functools.cache +def _load_tuned_configs(N: int, K: int) -> Optional[Dict[int, Dict[str, int]]]: + path = _CONFIGS_DIR / config_file_name(N, K) + if path.exists(): + logger.info("Using tuned GEMM configurations from %s.", path) + with open(path) as f: + return {int(m): config for m, config in json.load(f).items()} + logger.warning( + "No tuned GEMM configuration found for this device and shape (%s). " + "Falling back to heuristic defaults, performance may be suboptimal. " + "Generate a configuration with `tune_gemm(N=%d, K=%d)`.", + path.name, + N, + K, + ) + return None + + +def get_config(M: int, N: int, K: int) -> Dict[str, int]: + tuned = _load_tuned_configs(N, K) + if tuned is not None: + # Tuned Ms are spaced logarithmically, so pick the nearest in log space. + nearest_m = min(tuned, key=lambda m: abs(math.log(M / m))) + return tuned[nearest_m] + return default_config(M, N, K) +``` + +The configuration is then passed to the Triton kernel as its `constexpr` +and launch parameters (see +[`gemm.py`](https://github.com/huggingface/kernels/blob/main/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/gemm.py) +in the example): + +```python +def launch_gemm_kernel(a, b, out, config): + M, K = a.shape + N = b.shape[1] + grid = ( + triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv(N, config["BLOCK_SIZE_N"]), + ) + _gemm_kernel[grid](a, b, out, M, N, K, ..., **config) +``` + +## Generating the configurations + +The autotuner itself is ordinary benchmarking code: for each `M`, benchmark +every candidate configuration with `triton.testing.do_bench` and keep the +fastest. The example ships the tuner as part of the kernel (the `tune_gemm` +function in +[`tuning.py`](https://github.com/huggingface/kernels/blob/main/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/tuning.py)), +so users can also generate configurations for GPUs that the kernel author +did not tune. Candidates that do not fit the device (Triton raises +`OutOfResources`) are skipped. + +The example repository contains a small script, +[`tune.py`](https://github.com/huggingface/kernels/blob/main/examples/kernels/gemm-triton-autotune/tune.py), +that runs the tuner and writes the configuration files to the source tree: + +```bash +$ python tune.py --n 4096 --k 4096 +$ python tune.py --n 14336 --k 4096 +``` + +Commit the generated files, rebuild, and the configurations ship with the +kernel. When tuning a kernel that is not yet published, build it locally +(see [Develop locally](local-dev.md)) and point `LOCAL_KERNELS` at the +build: + +```bash +$ LOCAL_KERNELS=kernels-test/gemm-triton-autotune=build python tune.py --n 4096 --k 4096 +``` + +## Impact + +Tuned configurations are cheap to ship and can make a large difference. On +an NVIDIA L4, the tuned configuration for a `(1024, 4096) @ (4096, 4096)` +float16 GEMM is ~1.4× faster than the example's heuristic default (0.50 ms +vs. 0.71 ms) — on par with cuBLAS for this shape. diff --git a/docs/source/builder/writing-kernels.md b/docs/source/builder/writing-kernels.md index 905ba7ed..30679def 100644 --- a/docs/source/builder/writing-kernels.md +++ b/docs/source/builder/writing-kernels.md @@ -247,7 +247,9 @@ following options: - `src` (required): a list of source files and headers. - `pyext` (optional): the list of extensions for Python files. Default: - `["py", "pyi"]`. + `["py", "pyi"]`. Additional extensions can be listed to ship data files + with the kernel, such as `json` for + [Triton autotune configurations](triton-autotune.md). - `include` (optional): include directories relative to the project root. Default: `[]`. - `maxver` (optional): only build for this Torch version and earlier. Use cautiously, since this option produces @@ -297,7 +299,9 @@ the supported archs. These are then exported to `metadata.json` for consumption by e.g. the Hugging Face Hub. - `pyext` (optional): the list of extensions for Python files. Default: - `["py", "pyi"]`. + `["py", "pyi"]`. Additional extensions can be listed to ship data files + with the kernel, such as `json` for + [Triton autotune configurations](triton-autotune.md). - `cuda-capabilities` (optional): a list of CUDA compute capabilities the kernel supports (e.g. `["9.0", "10.0"]`). - `rocm-archs` (optional): a list of ROCm architectures the kernel supports diff --git a/examples/kernels/flake.nix b/examples/kernels/flake.nix index 2af64fa6..11ff0056 100644 --- a/examples/kernels/flake.nix +++ b/examples/kernels/flake.nix @@ -179,6 +179,11 @@ path = ./relu-triton; drv = sys: out: out.packages.${sys}.redistributable.torch-cuda; } + { + name = "gemm-triton-autotune-kernel"; + path = ./gemm-triton-autotune; + drv = sys: out: out.packages.${sys}.redistributable.torch-cuda; + } ]; # ROCm kernels to build in CI. @@ -188,6 +193,11 @@ path = ./relu-triton; drv = sys: out: out.packages.${sys}.redistributable.torch-rocm; } + { + name = "gemm-triton-autotune-kernel"; + path = ./gemm-triton-autotune; + drv = sys: out: out.packages.${sys}.redistributable.torch-rocm; + } { name = "relu-invalid-capability"; path = ./relu-invalid-capability; @@ -269,6 +279,11 @@ path = ./relu-triton; drv = sys: out: out.packages.${sys}.redistributable.torch-xpu; } + { + name = "gemm-triton-autotune-kernel"; + path = ./gemm-triton-autotune; + drv = sys: out: out.packages.${sys}.redistributable.torch-xpu; + } { name = "relu-kernel"; path = ./relu; diff --git a/examples/kernels/gemm-triton-autotune/CARD.md b/examples/kernels/gemm-triton-autotune/CARD.md new file mode 100644 index 00000000..b970ce41 --- /dev/null +++ b/examples/kernels/gemm-triton-autotune/CARD.md @@ -0,0 +1,60 @@ +--- +library_name: kernels +{% if license %}license: {{ license }} +{% endif %}--- + +This is the repository card of {{ repo_id }} that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated. + +## How to use +{% if functions %} + +```python +# make sure `kernels` is installed: `pip install -U kernels` +from kernels import get_kernel + +# If the org / user isn't a trusted publisher, pass `trust_remote_code=True` to the +# `get_kernel` call. You can find whether this kernel is from a trusted publisher +# by going to the kernel's Hub page and finding the "Trusted publisher" status at +# the top of the page. +kernel_module = get_kernel("{{ repo_id }}", version={{ version }}) +{{ functions[0] }} = kernel_module.{{ functions[0] }} + +{{ functions[0] }}(...) +``` +{% else %} + +Usage example not available. +{% endif %} + +## Available functions +{% if functions %} +{% for func in functions %} +- `{{ func }}` +{% endfor %} +{% else %} + +Function list not available. +{% endif %} +{% if layers %} + +## Available layers +{% for layer in layers %} +- `{{ layer }}` +{% endfor %} +{% endif %} + +## Benchmarks +{% if has_benchmark %} + +Benchmarking script is available for this kernel. Run `kernels benchmark {{ repo_id }} --version {{ version }}`. +{% else %} + +No benchmark available yet. +{% endif %} +{% if upstream %} + +## Source code + +Source code of this kernel originally comes from {{ upstream }} and it was repurposed for compatibility with `kernels`. +{% endif %} + diff --git a/examples/kernels/gemm-triton-autotune/build.toml b/examples/kernels/gemm-triton-autotune/build.toml new file mode 100644 index 00000000..39f79bc4 --- /dev/null +++ b/examples/kernels/gemm-triton-autotune/build.toml @@ -0,0 +1,19 @@ +[general] +name = "gemm-triton-autotune" +version = 1 +edition = 5 +license = "Apache-2.0" +backends = [ + "cuda", + "rocm", + "xpu", +] + +[general.hub] +repo-id = "kernels-test/gemm-triton-autotune" + +[torch-noarch] +pyext = [ + "json", + "py", +] diff --git a/examples/kernels/gemm-triton-autotune/flake.nix b/examples/kernels/gemm-triton-autotune/flake.nix new file mode 100644 index 00000000..29c0cc8f --- /dev/null +++ b/examples/kernels/gemm-triton-autotune/flake.nix @@ -0,0 +1,17 @@ +{ + description = "Flake for kernels tests"; + + inputs = { + kernel-builder.url = "path:../../.."; + }; + + outputs = + { + self, + kernel-builder, + }: + kernel-builder.lib.genKernelFlakeOutputs { + inherit self; + path = ./.; + }; +} diff --git a/examples/kernels/gemm-triton-autotune/tests/__init__.py b/examples/kernels/gemm-triton-autotune/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/kernels/gemm-triton-autotune/tests/conftest.py b/examples/kernels/gemm-triton-autotune/tests/conftest.py new file mode 100644 index 00000000..103c610f --- /dev/null +++ b/examples/kernels/gemm-triton-autotune/tests/conftest.py @@ -0,0 +1,12 @@ +import pytest +import torch + + +@pytest.fixture(scope="session") +def device() -> torch.device: + if hasattr(torch, "xpu") and torch.xpu.is_available(): + return torch.device("xpu") + elif torch.cuda.is_available(): + return torch.device("cuda") + else: + pytest.skip("No GPU available for Triton tests") diff --git a/examples/kernels/gemm-triton-autotune/tests/test_gemm.py b/examples/kernels/gemm-triton-autotune/tests/test_gemm.py new file mode 100644 index 00000000..d6c7ed2c --- /dev/null +++ b/examples/kernels/gemm-triton-autotune/tests/test_gemm.py @@ -0,0 +1,96 @@ +import json + +import kernels +import pytest +import torch + +gemm_kernel = kernels.get_kernel("kernels-test/gemm-triton-autotune", version=1) + +# GEMM output values grow with sqrt(K) for randn inputs, so tolerances are +# scaled by the magnitude of the output values. +DTYPE_TOLERANCES = { + torch.float16: {"rtol": 2e-2, "atol": 5e-1}, + torch.bfloat16: {"rtol": 4e-2, "atol": 4.0}, +} + + +@pytest.mark.kernels_ci +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("M", [1, 17, 256]) +@pytest.mark.parametrize("N,K", [(4096, 4096), (512, 511)]) +def test_gemm(device, dtype, M, N, K): + a = torch.randn(M, K, device=device, dtype=dtype) + b = torch.randn(K, N, device=device, dtype=dtype) + ref = (a.float() @ b.float()).to(dtype) + torch.testing.assert_close(gemm_kernel.gemm(a, b), ref, **DTYPE_TOLERANCES[dtype]) + + out = torch.empty(M, N, device=device, dtype=dtype) + gemm_kernel.gemm(a, b, out=out) + torch.testing.assert_close(out, ref, **DTYPE_TOLERANCES[dtype]) + + +@pytest.mark.kernels_ci +def test_gemm_validation(device): + a = torch.randn(4, 8, device=device, dtype=torch.float16) + with pytest.raises(ValueError, match="Incompatible"): + gemm_kernel.gemm(a, torch.randn(4, 8, device=device, dtype=torch.float16)) + with pytest.raises(ValueError, match="dtype"): + gemm_kernel.gemm(a, torch.randn(8, 4, device=device, dtype=torch.bfloat16)) + with pytest.raises(ValueError, match="dtype"): + gemm_kernel.gemm(a.float(), torch.randn(8, 4, device=device)) + with pytest.raises(ValueError, match="2D"): + gemm_kernel.gemm(a.flatten(), a.flatten()) + + +@pytest.mark.kernels_ci +def test_nearest_tuned_m_is_used(device, monkeypatch, tmp_path): + tuning = gemm_kernel.tuning + + small = tuning.default_config(16, 128, 128) + large = tuning.default_config(1024, 128, 128) + assert small != large + config_path = tmp_path / tuning.config_file_name(128, 128) + config_path.write_text(json.dumps({"16": small, "1024": large})) + + monkeypatch.setattr(tuning, "_CONFIGS_DIR", tmp_path) + tuning._load_tuned_configs.cache_clear() + try: + assert tuning.get_config(1, 128, 128) == small + assert tuning.get_config(32, 128, 128) == small + assert tuning.get_config(500, 128, 128) == large + assert tuning.get_config(100000, 128, 128) == large + finally: + tuning._load_tuned_configs.cache_clear() + + +@pytest.mark.kernels_ci +def test_fallback_to_default_config(device): + tuning = gemm_kernel.tuning + # No configuration is shipped for this shape, so the heuristic default + # should be used. + N, K = 123, 321 + assert tuning.get_config(64, N, K) == tuning.default_config(64, N, K) + + +@pytest.mark.kernels_ci +def test_shipped_config_is_used(device): + tuning = gemm_kernel.tuning + # Configurations tuned on an L4 GPU ship with the kernel (which is also + # the GPU that CI runs on). + if tuning.device_name() != "NVIDIA_L4": + pytest.skip("Shipped configurations were tuned for NVIDIA L4") + for N, K in [(4096, 4096), (14336, 4096)]: + assert tuning._load_tuned_configs(N, K) is not None + + +@pytest.mark.kernels_ci +def test_tune_gemm(device, tmp_path): + tuning = gemm_kernel.tuning + candidates = [tuning.default_config(16, 256, 256), tuning.default_config(64, 256, 256)] + path = gemm_kernel.tune_gemm(N=256, K=256, Ms=(1, 64), save_dir=tmp_path, candidates=candidates) + + assert path.parent == tmp_path + configs = json.loads(path.read_text()) + assert set(configs.keys()) == {"1", "64"} + for config in configs.values(): + assert config in candidates diff --git a/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/__init__.py b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/__init__.py new file mode 100644 index 00000000..4e239a0c --- /dev/null +++ b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/__init__.py @@ -0,0 +1,26 @@ +from typing import Optional + +import torch + +from ._ops import ops +from .gemm import _gemm # noqa: F401 — imported to register the op. +from .tuning import tune_gemm + +_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16) + + +def gemm(a: torch.Tensor, b: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor: + """Compute ``a @ b`` using a Triton GEMM kernel with tuned configurations.""" + if a.ndim != 2 or b.ndim != 2: + raise ValueError("gemm expects 2D tensors") + if a.shape[1] != b.shape[0]: + raise ValueError(f"Incompatible GEMM shapes: {tuple(a.shape)} @ {tuple(b.shape)}") + if a.dtype not in _SUPPORTED_DTYPES or a.dtype != b.dtype: + raise ValueError(f"gemm expects two {_SUPPORTED_DTYPES} tensors of the same dtype") + if out is None: + out = torch.empty(a.shape[0], b.shape[1], device=a.device, dtype=a.dtype) + ops.gemm(out, a, b) + return out + + +__all__ = ["gemm", "tune_gemm"] diff --git a/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/configs/N=14336,K=4096,device_name=NVIDIA_L4.json b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/configs/N=14336,K=4096,device_name=NVIDIA_L4.json new file mode 100644 index 00000000..45a90a97 --- /dev/null +++ b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/configs/N=14336,K=4096,device_name=NVIDIA_L4.json @@ -0,0 +1,50 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/configs/N=4096,K=4096,device_name=NVIDIA_L4.json b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/configs/N=4096,K=4096,device_name=NVIDIA_L4.json new file mode 100644 index 00000000..6fcaafa3 --- /dev/null +++ b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/configs/N=4096,K=4096,device_name=NVIDIA_L4.json @@ -0,0 +1,50 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/gemm.py b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/gemm.py new file mode 100644 index 00000000..1b7f728a --- /dev/null +++ b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/gemm.py @@ -0,0 +1,103 @@ +from typing import Dict + +import torch +import triton +import triton.language as tl + +from ._ops import add_op_namespace_prefix +from .tuning import get_config + + +@triton.jit +def _gemm_kernel( + a_ptr, + b_ptr, + c_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + acc = tl.dot(a, b, acc) + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + c = acc.to(c_ptr.dtype.element_ty) + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +def launch_gemm_kernel( + a: torch.Tensor, b: torch.Tensor, out: torch.Tensor, config: Dict[str, int] +) -> None: + """Launch the GEMM kernel with an explicit configuration. + + The tuner uses this to benchmark candidate configurations; regular calls + go through the ``gemm`` op, which looks up the configuration itself. + """ + M, K = a.shape + N = b.shape[1] + grid = ( + triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv(N, config["BLOCK_SIZE_N"]), + ) + _gemm_kernel[grid]( + a, + b, + out, + M, + N, + K, + a.stride(0), + a.stride(1), + b.stride(0), + b.stride(1), + out.stride(0), + out.stride(1), + **config, + ) + + +@torch.library.custom_op(add_op_namespace_prefix("gemm"), mutates_args={"out"}) +def _gemm(out: torch.Tensor, a: torch.Tensor, b: torch.Tensor) -> None: + M, K = a.shape + N = b.shape[1] + if M == 0 or N == 0: + return + launch_gemm_kernel(a, b, out, get_config(M, N, K)) + + +@_gemm.register_fake +def _(out: torch.Tensor, a: torch.Tensor, b: torch.Tensor) -> None: + pass diff --git a/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/tuning.py b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/tuning.py new file mode 100644 index 00000000..941a6f03 --- /dev/null +++ b/examples/kernels/gemm-triton-autotune/torch-ext/gemm_triton_autotune/tuning.py @@ -0,0 +1,189 @@ +"""Loading and generating tuned GEMM configurations. + +Tuned configurations are stored as JSON files in the ``configs`` directory +of this package, one file per ``(N, K, device)`` combination. Each file maps +an ``M`` value (the dimension that is typically only known at runtime, e.g. +the number of tokens) to the Triton configuration that performed best for +that ``M``: + + { + "16": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, ..., "num_warps": 4}, + "1024": {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, ...}, + ... + } + +At kernel launch, the configuration with the nearest tuned ``M`` is used. +When no configuration file exists for the current device and shape, a +heuristic default is used instead and a warning is logged once. +""" + +import functools +import json +import logging +import math +from pathlib import Path +from typing import Dict, Iterable, Optional, Sequence + +import torch +import triton + +logger = logging.getLogger(__name__) + +_CONFIGS_DIR = Path(__file__).parent / "configs" + +# The M values that are tuned by default: GEMMs are often skinny (a few +# tokens during decoding) or wide (large batches during prefill), so we +# cover a logarithmic range. +DEFAULT_TUNE_MS = (1, 16, 64, 256, 1024, 4096) + + +def device_name() -> str: + """Name of the current accelerator, as used in configuration file names.""" + if hasattr(torch, "xpu") and torch.xpu.is_available(): + name = torch.xpu.get_device_name() + else: + name = torch.cuda.get_device_name() + return name.replace(" ", "_") + + +def config_file_name(N: int, K: int) -> str: + return f"N={N},K={K},device_name={device_name()}.json" + + +@functools.lru_cache +def _load_tuned_configs(N: int, K: int) -> Optional[Dict[int, Dict[str, int]]]: + path = _CONFIGS_DIR / config_file_name(N, K) + if path.exists(): + logger.info("Using tuned GEMM configurations from %s.", path) + with open(path) as f: + return {int(m): config for m, config in json.load(f).items()} + logger.warning( + "No tuned GEMM configuration found for this device and shape (%s). " + "Falling back to heuristic defaults, performance may be suboptimal. " + "Generate a configuration with `tune_gemm(N=%d, K=%d)`.", + path.name, + N, + K, + ) + return None + + +def get_config(M: int, N: int, K: int) -> Dict[str, int]: + """Return the best known configuration for a ``(M, N, K)`` GEMM.""" + tuned = _load_tuned_configs(N, K) + if tuned: + # Pick the configuration tuned for the M closest to the runtime M + # (in log space, since tuned Ms are spaced logarithmically). + nearest_m = min(tuned, key=lambda m: abs(math.log(M / m))) + return tuned[nearest_m] + return default_config(M, N, K) + + +def default_config(M: int, N: int, K: int) -> Dict[str, int]: + """Heuristic configuration for GEMMs without tuned configurations. + + Deliberately conservative so that it runs on every supported device. + """ + return { + "BLOCK_SIZE_M": 16 if M <= 16 else 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + } + + +def candidate_configs() -> Iterable[Dict[str, int]]: + """The configuration search space used by ``tune_gemm``. + + Candidates that do not fit the device (e.g. exceed shared memory) are + skipped during tuning. + """ + for block_m in (16, 32, 64, 128): + for block_n in (32, 64, 128): + for block_k in (32, 64): + for num_warps in (4, 8): + # Wide warps only help for large tiles. + if num_warps == 8 and block_m * block_n < 64 * 128: + continue + for num_stages in (2, 3, 4): + yield { + "BLOCK_SIZE_M": block_m, + "BLOCK_SIZE_N": block_n, + "BLOCK_SIZE_K": block_k, + "GROUP_SIZE_M": 8, + "num_warps": num_warps, + "num_stages": num_stages, + } + + +def tune_gemm( + N: int, + K: int, + Ms: Sequence[int] = DEFAULT_TUNE_MS, + dtype: torch.dtype = torch.float16, + save_dir: Optional[Path] = None, + candidates: Optional[Sequence[Dict[str, int]]] = None, +) -> Path: + """Benchmark candidate configurations and save the best ones. + + For each ``M`` in ``Ms``, every candidate configuration is benchmarked + for an ``(M, K) @ (K, N)`` GEMM and the fastest is recorded. The result + is written as a JSON file named after ``(N, K)`` and the current device. + + ``save_dir`` defaults to this package's ``configs`` directory, so that + subsequent ``gemm`` calls in the same environment pick up the tuned + configurations. Kernel authors should pass the ``configs`` directory of + the kernel *source tree* instead (see ``tune.py`` in the repository + root) and commit the result, so that the configurations ship with the + kernel. + """ + # Import here to avoid a circular import at module load time. + from .gemm import launch_gemm_kernel + + device = "xpu" if hasattr(torch, "xpu") and torch.xpu.is_available() else "cuda" + if candidates is None: + candidates = list(candidate_configs()) + + best_configs: Dict[int, Dict[str, int]] = {} + for M in Ms: + a = torch.randn(M, K, device=device, dtype=dtype) + b = torch.randn(K, N, device=device, dtype=dtype) + out = torch.empty(M, N, device=device, dtype=dtype) + + best_time = math.inf + for config in candidates: + try: + time = triton.testing.do_bench( + lambda: launch_gemm_kernel(a, b, out, config) + ) + except triton.runtime.errors.OutOfResources: + # Configuration does not fit this device, skip it. + continue + if time < best_time: + best_time, best_configs[M] = time, config + + if M not in best_configs: + raise RuntimeError(f"No candidate configuration fits M={M}, N={N}, K={K}") + logger.info( + "Best configuration for M=%d, N=%d, K=%d: %s (%.4f ms)", + M, + N, + K, + best_configs[M], + best_time, + ) + + if save_dir is None: + save_dir = _CONFIGS_DIR + save_dir.mkdir(parents=True, exist_ok=True) + path = save_dir / config_file_name(N, K) + with open(path, "w") as f: + json.dump({str(m): config for m, config in sorted(best_configs.items())}, f, indent=4) + f.write("\n") + + # Make sure that new configurations are picked up by subsequent calls. + _load_tuned_configs.cache_clear() + + return path diff --git a/examples/kernels/gemm-triton-autotune/tune.py b/examples/kernels/gemm-triton-autotune/tune.py new file mode 100644 index 00000000..141f8212 --- /dev/null +++ b/examples/kernels/gemm-triton-autotune/tune.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python3 +"""Autotune the GEMM kernel and save the configurations to the source tree. + +The tuned configurations are written to +``torch-ext/gemm_triton_autotune/configs/``, named after the GEMM shape and +the current accelerator. Commit these files so that they ship with the +kernel. + +Tune the kernel as published on the Hub: + + python tune.py --n 4096 --k 4096 + +Or tune a local build (see the "Develop locally" chapter of the +kernel-builder documentation for building into ``build/``): + + LOCAL_KERNELS=kernels-test/gemm-triton-autotune=build python tune.py --n 4096 --k 4096 +""" + +import argparse +import logging +from pathlib import Path + +import torch + +import kernels + +SAVE_DIR = Path(__file__).parent / "torch-ext" / "gemm_triton_autotune" / "configs" + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--n", type=int, required=True, help="N dimension of the GEMM") + parser.add_argument("--k", type=int, required=True, help="K dimension of the GEMM") + parser.add_argument( + "--m", + type=int, + nargs="+", + default=None, + help="M values to tune (default: a logarithmic range)", + ) + parser.add_argument("--dtype", choices=["float16", "bfloat16"], default="float16") + parser.add_argument( + "--save-dir", + type=Path, + default=SAVE_DIR, + help="Directory to write the configuration file to", + ) + args = parser.parse_args() + + logging.basicConfig(level=logging.INFO) + + gemm_kernel = kernels.get_kernel("kernels-test/gemm-triton-autotune", version=1) + + tune_kwargs = {} + if args.m is not None: + tune_kwargs["Ms"] = tuple(args.m) + + path = gemm_kernel.tune_gemm( + N=args.n, + K=args.k, + dtype=getattr(torch, args.dtype), + save_dir=args.save_dir, + **tune_kwargs, + ) + print(f"Configurations written to {path}") + + +if __name__ == "__main__": + main() diff --git a/nix-builder/tests/Dockerfile.test-kernel b/nix-builder/tests/Dockerfile.test-kernel index 2a605290..7329b5cb 100644 --- a/nix-builder/tests/Dockerfile.test-kernel +++ b/nix-builder/tests/Dockerfile.test-kernel @@ -77,6 +77,7 @@ COPY relu-kernel-cpu ./relu-kernel-cpu COPY cutlass-gemm-kernel ./cutlass-gemm-kernel COPY cutlass-gemm-tvm-ffi-kernel ./cutlass-gemm-tvm-ffi-kernel COPY relu-triton-kernel ./relu-triton-kernel +COPY gemm-triton-autotune-kernel ./gemm-triton-autotune-kernel COPY silu-and-mul-kernel ./silu-and-mul-kernel COPY extra-data ./extra-data COPY cpp20-symbols-kernel ./cpp20-symbols-kernel @@ -87,6 +88,7 @@ COPY examples/kernels/cutlass-gemm/tests ./cutlass_gemm_tests COPY examples/kernels/cutlass-gemm-tvm-ffi/tests ./cutlass_gemm_tvm_ffi_tests COPY examples/kernels/cpp20-symbols/tests ./cpp20_symbols_tests COPY examples/kernels/relu-triton/tests ./relu_triton_tests +COPY examples/kernels/gemm-triton-autotune/tests ./gemm_triton_autotune_tests # Run tests ADD nix-builder/tests/run-tests.sh ./run-tests.sh diff --git a/nix-builder/tests/run-tests.sh b/nix-builder/tests/run-tests.sh index 33426b81..7374aba1 100644 --- a/nix-builder/tests/run-tests.sh +++ b/nix-builder/tests/run-tests.sh @@ -9,6 +9,7 @@ RELU_TVM_FFI_PATH=$(echo relu-tvm-ffi-kernel/tvm-ffi*) CUTLASS_PATH=$(echo cutlass-gemm-kernel/torch*) CUTLASS_TVM_FFI_PATH=$(echo cutlass-gemm-tvm-ffi-kernel/tvm-ffi*) RELU_TRITON_PATH=$(echo relu-triton-kernel/torch*) +GEMM_TRITON_AUTOTUNE_PATH=$(echo gemm-triton-autotune-kernel/torch*) SILU_MUL_PATH=$(echo silu-and-mul-kernel/torch*) RELU_CPU_PATH=$(echo relu-kernel-cpu/torch*) CPP20_SYMBOLS_PATH=$(echo cpp20-symbols-kernel/torch*) @@ -19,6 +20,9 @@ LOCAL_KERNELS="kernels-test/extra-data=${EXTRA_DATA_PATH}:kernels-test/relu=${RE LOCAL_KERNELS="kernels-test/relu-triton=${RELU_TRITON_PATH}" \ .venv/bin/pytest relu_triton_tests +LOCAL_KERNELS="kernels-test/gemm-triton-autotune=${GEMM_TRITON_AUTOTUNE_PATH}" \ + .venv/bin/pytest gemm_triton_autotune_tests + # We only care about importing, the kernel is trivial. LOCAL_KERNELS="kernels-test/silu-and-mul=${SILU_MUL_PATH}" \ .venv/bin/python -c "import kernels; kernels.get_kernel('kernels-test/silu-and-mul', version=1)"