ocnn 2.2.7__tar.gz → 2.3.0__tar.gz

{ocnn-2.2.7/ocnn.egg-info → ocnn-2.3.0}/PKG-INFO

@@ -1,13 +1,12 @@
 Metadata-Version: 2.4
 Name: ocnn
-Version: 2.2.7
+Version: 2.3.0
 Summary: Octree-based Sparse Convolutional Neural Networks
 Home-page: https://github.com/octree-nn/ocnn-pytorch
 Author: Peng-Shuai Wang
 Author-email: wangps@hotmail.com
 License: MIT
 Classifier: Programming Language :: Python :: 3
-Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
 Requires-Python: >=3.6
 Description-Content-Type: text/markdown
@@ -84,6 +83,13 @@ octrees to perform convolution operations. Of course, it also supports other 3D
 data formats, such as meshes and volumetric grids, which can be converted into
 octrees to leverage the library's capabilities.
 
+## Updates
+
+- **2026.02.02**: Release `v2.3.0`, incorporating Triton to accelerate
+  octree-based sparse convolution in the upcoming release. OctreeConv is even
+  **2.5 times faster than the latest spconv**!
+- **2025.12.18**: Release `v2.2.8`, improving neighbor search efficiency.
+
 
 ## Key benefits of ocnn-pytorch
 
@@ -93,10 +99,9 @@ octrees to leverage the library's capabilities.
   configure the compiling environment.
 
 - **Efficiency**. The ocnn-pytorch is very efficient compared with other sparse
-  convolution frameworks.  It only takes 18 hours to train the network on
-  ScanNet for 600 epochs with 4 V100 GPUs. For reference, under the same
-  training settings, MinkowskiNet 0.4.3 takes 60 hours and MinkowskiNet 0.5.4
-  takes 30 hours.
+  convolution frameworks. It is **even 2.5 times faster than the latest spconv
+  implementation**! Check the benchmark [code](test/benchmark_conv.py) and
+  [results](test/benchmark/results.png) for details. ✨
 
 ## Citation
 

{ocnn-2.2.7 → ocnn-2.3.0}/README.md

@@ -54,6 +54,13 @@ octrees to perform convolution operations. Of course, it also supports other 3D
 data formats, such as meshes and volumetric grids, which can be converted into
 octrees to leverage the library's capabilities.
 
+## Updates
+
+- **2026.02.02**: Release `v2.3.0`, incorporating Triton to accelerate
+  octree-based sparse convolution in the upcoming release. OctreeConv is even
+  **2.5 times faster than the latest spconv**!
+- **2025.12.18**: Release `v2.2.8`, improving neighbor search efficiency.
+
 
 ## Key benefits of ocnn-pytorch
 
@@ -63,10 +70,9 @@ octrees to leverage the library's capabilities.
   configure the compiling environment.
 
 - **Efficiency**. The ocnn-pytorch is very efficient compared with other sparse
-  convolution frameworks.  It only takes 18 hours to train the network on
-  ScanNet for 600 epochs with 4 V100 GPUs. For reference, under the same
-  training settings, MinkowskiNet 0.4.3 takes 60 hours and MinkowskiNet 0.5.4
-  takes 30 hours.
+  convolution frameworks. It is **even 2.5 times faster than the latest spconv
+  implementation**! Check the benchmark [code](test/benchmark_conv.py) and
+  [results](test/benchmark/results.png) for details. ✨
 
 ## Citation
 

{ocnn-2.2.7 → ocnn-2.3.0}/ocnn/__init__.py

@@ -12,7 +12,7 @@ from . import models
 from . import dataset
 from . import utils
 
-__version__ = '2.2.7'
+__version__ = '2.3.0'
 
 __all__ = [
     'octree',

{ocnn-2.2.7 → ocnn-2.3.0}/ocnn/models/resnet.py

@@ -15,7 +15,7 @@ class ResNet(torch.nn.Module):
   '''
 
   def __init__(self, in_channels: int, out_channels: int, resblock_num: int,
-               stages: int, nempty: bool = False):
+               stages: int, nempty: bool = False, dropout: float = 0.5):
     super().__init__()
     self.in_channels = in_channels
     self.out_channels = out_channels
@@ -36,7 +36,7 @@ class ResNet(torch.nn.Module):
     # self.header = torch.nn.Linear(channels[-1], out_channels, bias=True)
     self.header = torch.nn.Sequential(
         ocnn.modules.FcBnRelu(channels[-1], 512),
-        torch.nn.Dropout(p=0.5),
+        torch.nn.Dropout(p=dropout),
         torch.nn.Linear(512, out_channels))
 
   def forward(self, data: torch.Tensor, octree: Octree, depth: int):

{ocnn-2.2.7 → ocnn-2.3.0}/ocnn/nn/__init__.py

@@ -21,7 +21,7 @@ from .octree_norm import (OctreeBatchNorm, OctreeGroupNorm,
                           OctreeInstanceNorm, OctreeNorm)
 from .octree_drop import OctreeDropPath
 from .octree_align import search_value, octree_align
-
+from .octree_conv_t import OctreeConvTriton, OctreeConvT, convert_conv_triton
 
 __all__ = [
     'octree2voxel',
@@ -39,6 +39,7 @@ __all__ = [
     'OctreeInstanceNorm', 'OctreeBatchNorm', 'OctreeGroupNorm', 'OctreeNorm',
     'OctreeDropPath',
     'search_value', 'octree_align',
+    'OctreeConvTriton', 'OctreeConvT', 'convert_conv_triton',
 ]
 
 classes = __all__

ocnn-2.3.0/ocnn/nn/kernels/__init__.py

@@ -0,0 +1,14 @@
+from .conv_fwd_implicit_gemm_splitk import conv_fwd_implicit_gemm_splitk
+from .conv_bwd_implicit_gemm_splitk import conv_bwd_implicit_gemm_splitk
+from .conv_bwd_implicit_gemm import conv_bwd_implicit_gemm
+from .conv_fwd_implicit_gemm import conv_fwd_implicit_gemm
+
+__all__ = [
+    'conv_fwd_implicit_gemm_splitk',
+    'conv_bwd_implicit_gemm_splitk',
+    'conv_bwd_implicit_gemm',
+    'conv_fwd_implicit_gemm',
+]
+
+from .autotuner import load_autotune_cache
+load_autotune_cache()

ocnn-2.3.0/ocnn/nn/kernels/autotuner.py

@@ -0,0 +1,416 @@
+import builtins
+import os
+import json
+import importlib
+import pkgutil
+import torch
+import triton
+import time
+import inspect
+from filelock import FileLock
+from typing import Dict, Mapping
+
+VERBOSE_AUTOTUNE = os.getenv('TRITON_PRINT_AUTOTUNING', '0') == '1'
+AUTOSAVE_AUTOTUNE_CACHE = os.getenv('OCNN_AUTOSAVE_AUTOTUNE', '1') == '1'
+AUTOTUNE_CACHE_PATH = os.getenv('OCNN_AUTOTUNE_CACHE_PATH',
+                                os.path.expanduser('~/.ocnnconvt/autotune_cache.json'))
+
+
+class TritonPersistentCacheAutotuner(triton.runtime.Autotuner):
+    def __init__(
+        self,
+        fn,
+        arg_names,
+        configs,
+        key,
+        reset_to_zero,
+        restore_value,
+        pre_hook=None,
+        post_hook=None,
+        prune_configs_by: Dict = None,
+        warmup=None,
+        rep=None,
+        use_cuda_graph=False,
+        do_bench=None,
+    ):
+        super().__init__(
+            fn,
+            arg_names,
+            configs,
+            key,
+            reset_to_zero,
+            restore_value,
+            pre_hook,
+            post_hook,
+            prune_configs_by,
+            warmup,
+            rep,
+            use_cuda_graph,
+            do_bench,
+        )
+
+    def run(self, *args, **kwargs):
+        self.nargs = dict(zip(self.arg_names, args))
+        used_cached_result = True
+        if len(self.configs) > 1:
+            all_args = {**self.nargs, **kwargs}
+            _args = {k: v for (k, v) in all_args.items() if k in self.arg_names}
+            key = [_args[key] for key in self.keys if key in _args]
+            for _, arg in _args.items():
+                if hasattr(arg, "dtype"):
+                    key.append(str(arg.dtype))
+            key = str(tuple(key))
+            if key not in self.cache:
+                # prune configs
+                used_cached_result = False
+                pruned_configs = self.prune_configs(kwargs)
+                bench_start = time.time()
+                timings = {config: self._bench(*args, config=config, **kwargs) for config in pruned_configs}
+                bench_end = time.time()
+                self.bench_time = bench_end - bench_start
+                self.cache[key] = builtins.min(timings, key=timings.get)
+                full_nargs = {**self.nargs, **kwargs, **self.cache[key].all_kwargs()}
+                self.pre_hook(full_nargs, reset_only=True)
+                self.configs_timings = timings
+            config = self.cache[key]
+        else:
+            config = self.configs[0]
+        self.best_config = config
+        if os.getenv("TRITON_PRINT_AUTOTUNING", None) == "1" and not used_cached_result:
+            print(f"Triton autotuning for function {self.base_fn.__name__} finished after "
+                  f"{self.bench_time:.2f}s; best config selected: {self.best_config};")
+        if AUTOSAVE_AUTOTUNE_CACHE and not used_cached_result:
+            save_autotune_cache()
+        if config.pre_hook is not None:
+            full_nargs = {**self.nargs, **kwargs, **config.all_kwargs()}
+            config.pre_hook(full_nargs)
+        ret = self.fn.run(
+            *args,
+            **kwargs,
+            **config.all_kwargs(),
+        )
+        self.nargs = None
+        return ret
+
+    def prune_configs(self, kwargs):
+        pruned_configs = self.configs
+        if self.early_config_prune:
+            pruned_configs = self.early_config_prune(self.configs, self.nargs, **kwargs)
+        if self.perf_model:
+            top_k = self.configs_top_k
+            if isinstance(top_k, float) and top_k <= 1.0:
+                top_k = int(len(self.configs) * top_k)
+            if len(pruned_configs) > top_k:
+                est_timing = {
+                    config: self.perf_model(
+                        **self.nargs,
+                        **kwargs,
+                        **config.all_kwargs(),
+                    )
+                    for config in pruned_configs
+                }
+                pruned_configs = sorted(est_timing.keys(), key=lambda x: est_timing[x])[:top_k]
+        return pruned_configs
+
+    def warmup(self, *args, **kwargs):
+        self.nargs = dict(zip(self.arg_names, args))
+        ret = []
+        for config in self.prune_configs(kwargs):
+            ret.append(self.fn.warmup(
+                *args,
+                **kwargs,
+                **config.all_kwargs(),
+            ))
+        self.nargs = None
+        return ret
+
+
+def triton_autotune(configs, key, prune_configs_by=None, reset_to_zero=None, restore_value=None, pre_hook=None, post_hook=None,
+             warmup=None, rep=None, use_cuda_graph=False, do_bench=None):
+    """
+    Decorator for auto-tuning a :code:`triton.jit`'d function.
+
+    .. highlight:: python
+    .. code-block:: python
+
+        @triton_autotune(configs=[
+            triton.Config(kwargs={'BLOCK_SIZE': 128}, num_warps=4),
+            triton.Config(kwargs={'BLOCK_SIZE': 1024}, num_warps=8),
+          ],
+          key=['x_size'] # the two above configs will be evaluated anytime
+                         # the value of x_size changes
+        )
+        @triton.jit
+        def kernel(x_ptr, x_size, **META):
+            BLOCK_SIZE = META['BLOCK_SIZE']
+    :note: When all the configurations are evaluated, the kernel will run multiple times.
+           This means that whatever value the kernel updates will be updated multiple times.
+           To avoid this undesired behavior, you can use the `reset_to_zero` argument, which
+           resets the value of the provided tensor to `zero` before running any configuration.
+
+    If the environment variable :code:`TRITON_PRINT_AUTOTUNING` is set to
+    :code:`"1"`, Triton will print a message to stdout after autotuning each
+    kernel, including the time spent autotuning and the best configuration.
+
+    :param configs: a list of :code:`triton.Config` objects
+    :type configs: list[triton.Config]
+    :param key: a list of argument names whose change in value will trigger the evaluation of all provided configs.
+    :type key: list[str]
+    :param prune_configs_by: a dict of functions that are used to prune configs, fields:
+        'perf_model': performance model used to predicate running time with different configs, returns running time
+        'top_k': number of configs to bench
+        'early_config_prune'(optional): a function used to do early prune (eg, num_stages). It takes configs:List[Config] as its input, and returns pruned configs.
+    :param reset_to_zero: a list of argument names whose value will be reset to zero before evaluating any configs.
+    :type reset_to_zero: list[str]
+    :param restore_value: a list of argument names whose value will be restored after evaluating any configs.
+    :type restore_value: list[str]
+    :param pre_hook: a function that will be called before the kernel is called.
+        This overrides the default pre_hook used for 'reset_to_zero' and 'restore_value'.
+        'kwargs': a dict of all arguments passed to the kernel.
+        'reset_only': a boolean indicating whether the pre_hook is called to reset the values only, without a corresponding post_hook.
+    :type pre_hook: lambda args, reset_only
+    :param post_hook: a function that will be called after the kernel is called.
+        This overrides the default post_hook used for 'restore_value'.
+        'kwargs': a dict of all arguments passed to the kernel.
+        'exception': the exception raised by the kernel in case of a compilation or runtime error.
+    :type post_hook: lambda args, exception
+    :param warmup: warmup time (in ms) to pass to benchmarking (deprecated).
+    :type warmup: int
+    :param rep: repetition time (in ms) to pass to benchmarking (deprecated).
+    :type rep: int
+    :param do_bench: a benchmark function to measure the time of each run.
+    :type do_bench: lambda fn, quantiles
+    """
+
+    def decorator(fn):
+        return TritonPersistentCacheAutotuner(
+            fn, fn.arg_names, configs, key, reset_to_zero, restore_value, pre_hook=pre_hook,
+            post_hook=post_hook, prune_configs_by=prune_configs_by, warmup=warmup, rep=rep,
+            use_cuda_graph=use_cuda_graph
+        )
+
+    return decorator
+
+
+class PersistentCacheAutoTuner:
+    def __init__(
+        self,
+        kernel,
+        configs=None,
+        key=None,
+        config_fn=None,
+        key_fn=None,
+        warmup=3,
+        runs=10,
+        verbose=False,
+    ):
+        """
+        AutoTuner is a wrapper class for a kernel that automatically tunes the kernel parameters to achieve the best performance.
+
+        Args:
+            kernel: A callable object that takes in input arguments and returns the output.
+            configs: A list of Config objects that define the possible kernel parameters and their values.
+            key: A list of argument names that retune the kernel on change.
+            config_fn: A function that takes in the input arguments and returns configs to be used for autotuning.
+            key_fn: A function that takes in the input arguments and returns the key used to cache the tuning results.
+                    Once the key changes, the autotuning will be rerun.
+            warmup: The number of warmup runs to discard before measuring the execution time.
+            runs: The number of runs to measure the execution time.
+            verbose: Whether to print the autotuning results.
+        """
+        assert config_fn or configs, "Either configs or config_fn must be provided"
+        assert key_fn or key, "Either key or key_fn must be provided"
+        self.kernel = kernel
+        self.configs = configs
+        self.key = key
+        self.config_fn = config_fn
+        self.key_fn = key_fn
+        self.warmup = warmup
+        self.runs = runs
+        self.verbose = verbose
+        self.kernel_arg_names = inspect.getfullargspec(kernel).args
+        self.cache = {}
+
+    def _args_to_kwargs(self, args, kwargs):
+        # Convert args to kwargs
+        arg_names = self.kernel_arg_names
+        arg_dict = dict(zip(arg_names, args))
+        arg_dict.update(kwargs)
+        return arg_dict
+
+    def __call__(self, *args, **kwargs):
+        arg_dict = self._args_to_kwargs(args, kwargs)
+
+        # Determine key
+        key = self.key_fn(*args, **kwargs) if self.key_fn else tuple(arg_dict[k] for k in self.key)
+        key = str(key)
+
+        # If key changes, rerun autotune
+        used_cached_result = True
+        if key not in self.cache:
+            used_cached_result = False
+            if self.verbose:
+                print(f"Running autotuning for {self.kernel.__name__} with key {key}")
+            configs = self.configs if self.configs else self.config_fn(*args, **kwargs)
+            if self.verbose:
+                print(f"Configs: {configs}")
+            best_config = self._benchmark(args, kwargs, configs)
+            if self.verbose:
+                print(f"Best config for {self.kernel.__name__} with key {key}: {best_config}")
+            self.cache[key] = best_config
+        else:
+            if self.verbose:
+                print('Using cached config for {} with key {}'.format(self.kernel.__name__, key))
+                print('Config: {}'.format(self.cache[key]))
+
+        if AUTOSAVE_AUTOTUNE_CACHE and not used_cached_result:
+            save_autotune_cache()
+
+        # Run the kernel with the best config
+        return self.kernel(*args, **kwargs, **self.cache[key])
+
+    def _benchmark(self, args, kwargs, configs):
+        best_time = float('inf')
+        best_config = None
+
+        if len(configs) == 1:
+            best_config = configs[0]
+        else:
+            for config in configs:
+                # Run the kernel and measure execution time
+                for _ in range(self.warmup):
+                    self.kernel(*args, **kwargs, **config)
+                torch.cuda.synchronize()
+                start = time.time()
+                for _ in range(self.runs):
+                    self.kernel(*args, **kwargs, **config)
+                torch.cuda.synchronize()
+                elapsed = (time.time() - start) / self.runs
+                if self.verbose:
+                    print(f"Config {config}: {elapsed} seconds")
+                # Update the best config if the execution time is better
+                if elapsed < best_time:
+                    best_time = elapsed
+                    best_config = config
+
+        return best_config
+
+
+
+
+def autotune(
+    configs=None,
+    key=None,
+    config_fn=None,
+    key_fn=None,
+    warmup=3,
+    runs=10,
+    verbose=VERBOSE_AUTOTUNE
+):
+    def decorator(kernel):
+        return PersistentCacheAutoTuner(kernel, configs, key, config_fn, key_fn, warmup, runs, verbose)
+    return decorator
+
+
+def walk_package(package_name, fn):
+    try:
+        package = importlib.import_module(package_name)
+    except ModuleNotFoundError:
+        print(f"Package {package_name} not found.")
+        return
+
+    if not hasattr(package, '__path__'):
+        print(f"{package_name} is not a package.")
+        return
+
+    for _, module_name, is_pkg in pkgutil.iter_modules(package.__path__):
+        full_module_name = f"{package_name}.{module_name}"
+        if is_pkg:
+            walk_package(full_module_name, fn)
+        else:
+            fn(full_module_name)
+
+
+def get_autotune_cache():
+    cache = {}
+    device_name = torch.cuda.get_device_name()
+    if device_name not in cache:
+        cache[device_name] = {}
+
+    def save_cache(full_module_name):
+        module = importlib.import_module(full_module_name)
+        for attr_name, attr in module.__dict__.items():
+            cache_key = f"{full_module_name}.{attr_name}"
+            if isinstance(attr, PersistentCacheAutoTuner):
+                cache[device_name][cache_key] = attr.cache
+            elif isinstance(attr, TritonPersistentCacheAutotuner):
+                cache[device_name][cache_key] = {k: v.__dict__ for k, v in attr.cache.items()}
+
+    walk_package('ocnn.nn.kernels', save_cache)
+
+    return cache
+
+
+def save_autotune_cache(path=None):
+    path = path or AUTOTUNE_CACHE_PATH
+    lock_path = path + ".lock"
+
+    with FileLock(lock_path):
+        if os.path.exists(path):
+            with open(path, 'r') as f:
+                cache = json.load(f)
+        else:
+            cache = {}
+        # Merge existing cache with new cache
+        cache.update(get_autotune_cache())
+
+        tmp_path = path + ".tmp"
+        with open(tmp_path, 'w') as f:
+            json.dump(cache, f, indent=4)
+            f.flush()
+            os.fsync(f.fileno())
+        os.replace(tmp_path, path)
+
+
+def load_autotune_cache(path_or_cache=None):
+    cache = None
+
+    # Preserve path-based loading, but allow callers to provide a preloaded cache object.
+    if path_or_cache is None or isinstance(path_or_cache, (str, os.PathLike)):
+        path = path_or_cache or AUTOTUNE_CACHE_PATH
+        lock_path = path + ".lock"
+
+        if not os.path.exists(path):
+            return
+
+        with FileLock(lock_path):
+            with open(path, 'r') as f:
+                cache = json.load(f)
+    elif isinstance(path_or_cache, Mapping):
+        cache = path_or_cache
+    else:
+        raise TypeError("load_autotune_cache expects a path or a mapping")
+
+    if cache is None:
+        return
+
+    device_name = torch.cuda.get_device_name()
+    if device_name not in cache and "*" not in cache:
+        return
+    if "*" in cache and device_name not in cache:
+        device_name = "*"
+
+    def load_cache(full_module_name):
+        module = importlib.import_module(full_module_name)
+        for attr_name, attr in module.__dict__.items():
+            cache_key = f"{full_module_name}.{attr_name}"
+            if isinstance(attr, PersistentCacheAutoTuner):
+                if cache_key in cache[device_name]:
+                    attr.cache = cache[device_name][cache_key]
+            elif isinstance(attr, TritonPersistentCacheAutotuner):
+                if cache_key in cache[device_name]:
+                    for k, v in cache[device_name][cache_key].items():
+                        attr.cache[k] = triton.runtime.Config(None)
+                        attr.cache[k].__dict__.update(v)
+
+    walk_package('ocnn.nn.kernels', load_cache)

ocnn-2.3.0/ocnn/nn/kernels/config.py

@@ -0,0 +1,67 @@
+import os
+import triton
+from .utils import get_autotune_config
+
+allow_tf32 = os.getenv('OCNN_ALLOW_TF32', '1') == '1'
+
+autotune_config = get_autotune_config(
+    platform={
+        'cuda': [
+            triton.Config({'B1': 128, 'B2': 256, 'BK': 64}, num_stages=3, num_warps=8),
+            triton.Config({'B1': 64,  'B2': 256, 'BK': 32}, num_stages=4, num_warps=4),
+            triton.Config({'B1': 128, 'B2': 128, 'BK': 32}, num_stages=4, num_warps=4),
+            triton.Config({'B1': 128, 'B2': 64,  'BK': 32}, num_stages=4, num_warps=4),
+            triton.Config({'B1': 64,  'B2': 128, 'BK': 32}, num_stages=4, num_warps=4),
+            triton.Config({'B1': 128, 'B2': 32,  'BK': 32}, num_stages=4, num_warps=4),
+            triton.Config({'B1': 64,  'B2': 32,  'BK': 32}, num_stages=5, num_warps=2),
+            triton.Config({'B1': 32,  'B2': 64,  'BK': 32}, num_stages=5, num_warps=2),
+        ],
+        'hip': [
+            triton.Config({'B1': 128, 'B2': 256, 'BK': 16, 'waves_per_eu': 2}, num_warps=4, num_stages=2),
+            triton.Config({'B1': 256, 'B2': 256, 'BK': 16, 'waves_per_eu': 2}, num_warps=8, num_stages=2),
+            triton.Config({'B1': 128, 'B2': 128, 'BK': 32, 'waves_per_eu': 2}, num_warps=8, num_stages=2),
+            triton.Config({'B1': 64,  'B2': 128, 'BK': 32, 'waves_per_eu': 3}, num_warps=4, num_stages=2),
+            triton.Config({'B1': 64,  'B2': 64,  'BK': 32, 'waves_per_eu': 8}, num_warps=4, num_stages=2),
+        ]
+    },
+    device={
+        'A100': [
+            triton.Config({'B1': 256, 'B2': 128, 'BK': 64}, num_stages=4, num_warps=8),
+            triton.Config({'B1': 256, 'B2': 128, 'BK': 32}, num_stages=4, num_warps=8),
+            triton.Config({'B1': 128, 'B2': 256, 'BK': 64}, num_stages=4, num_warps=8),
+            triton.Config({'B1': 128, 'B2': 256, 'BK': 32}, num_stages=4, num_warps=8),
+            triton.Config({'B1': 256, 'B2': 64,  'BK': 64}, num_stages=4, num_warps=4),
+            triton.Config({'B1': 256, 'B2': 64,  'BK': 32}, num_stages=4, num_warps=4),
+            triton.Config({'B1': 64,  'B2': 256, 'BK': 64}, num_stages=4, num_warps=4),
+            triton.Config({'B1': 64,  'B2': 256, 'BK': 32}, num_stages=4, num_warps=4),
+            triton.Config({'B1': 128, 'B2': 128, 'BK': 64}, num_stages=4, num_warps=4),
+            triton.Config({'B1': 128, 'B2': 128, 'BK': 32}, num_stages=4, num_warps=4),
+            triton.Config({'B1': 128, 'B2': 64,  'BK': 32}, num_stages=4, num_warps=4),
+            triton.Config({'B1': 128, 'B2': 64,  'BK': 32}, num_stages=4, num_warps=2),
+            triton.Config({'B1': 64,  'B2': 128, 'BK': 32}, num_stages=4, num_warps=4),
+            triton.Config({'B1': 64,  'B2': 128, 'BK': 32}, num_stages=4, num_warps=2),
+            triton.Config({'B1': 64,  'B2': 64,  'BK': 64}, num_stages=4, num_warps=2),
+            triton.Config({'B1': 64,  'B2': 64,  'BK': 32}, num_stages=4, num_warps=2),
+        ],
+        'MI300X': [
+            triton.Config({'B1': 256, 'B2': 256, 'BK': 64, 'waves_per_eu': 2}, num_stages=2, num_warps=16),
+            triton.Config({'B1': 256, 'B2': 256, 'BK': 64, 'waves_per_eu': 2, 'kpack': 2, 'matrix_instr_nonkdim': 16}, num_stages=2, num_warps=8),
+            triton.Config({'B1': 256, 'B2': 128, 'BK': 64, 'waves_per_eu': 2}, num_stages=2, num_warps=16),
+            triton.Config({'B1': 256, 'B2': 128, 'BK': 64, 'waves_per_eu': 2, 'kpack': 2, 'matrix_instr_nonkdim': 16}, num_stages=2, num_warps=8),
+            triton.Config({'B1': 128, 'B2': 256, 'BK': 64, 'waves_per_eu': 2}, num_stages=2, num_warps=16),
+            triton.Config({'B1': 128, 'B2': 256, 'BK': 64, 'waves_per_eu': 2, 'kpack': 2, 'matrix_instr_nonkdim': 16}, num_stages=2, num_warps=8),
+            triton.Config({'B1': 256, 'B2': 64,  'BK': 32, 'waves_per_eu': 2}, num_stages=2, num_warps=8),
+            triton.Config({'B1': 256, 'B2': 64,  'BK': 32, 'waves_per_eu': 2, 'kpack': 2, 'matrix_instr_nonkdim': 16}, num_stages=2, num_warps=8),
+            triton.Config({'B1': 64,  'B2': 256, 'BK': 32, 'waves_per_eu': 2}, num_stages=2, num_warps=8),
+            triton.Config({'B1': 64,  'B2': 256, 'BK': 32, 'waves_per_eu': 2, 'kpack': 2, 'matrix_instr_nonkdim': 16}, num_stages=2, num_warps=8),
+            triton.Config({'B1': 128, 'B2': 128, 'BK': 64, 'waves_per_eu': 2}, num_stages=2, num_warps=8),
+            triton.Config({'B1': 128, 'B2': 128, 'BK': 64, 'waves_per_eu': 2, 'kpack': 2, 'matrix_instr_nonkdim': 16}, num_stages=2, num_warps=8),
+            triton.Config({'B1': 128, 'B2': 64,  'BK': 64, 'waves_per_eu': 2}, num_stages=2, num_warps=4),
+            triton.Config({'B1': 128, 'B2': 64,  'BK': 64, 'waves_per_eu': 2, 'kpack': 2, 'matrix_instr_nonkdim': 16}, num_stages=2, num_warps=4),
+            triton.Config({'B1': 64,  'B2': 128, 'BK': 64, 'waves_per_eu': 2}, num_stages=2, num_warps=4),
+            triton.Config({'B1': 64,  'B2': 128, 'BK': 64, 'waves_per_eu': 2, 'kpack': 2, 'matrix_instr_nonkdim': 16}, num_stages=2, num_warps=4),
+            triton.Config({'B1': 64,  'B2': 64,  'BK': 64, 'waves_per_eu': 2}, num_stages=2, num_warps=2),
+            triton.Config({'B1': 64,  'B2': 64,  'BK': 64, 'waves_per_eu': 2, 'kpack': 2, 'matrix_instr_nonkdim': 16}, num_stages=2, num_warps=2),
+        ],
+    }
+)