Some lints and refactor

deep_gemm/jit_kernels/gemm.py

@@ -3,8 +3,8 @@ import torch
 from functools import lru_cache
 from typing import Tuple
 
-from ..jit.utils import GemmType, make_2d_tma_a_desc, make_2d_tma_b_desc, make_2d_tma_d_desc, make_2d_tma_scales_a_desc
-
+from .runtime import FP8GemmRuntime, generate
+from .runtime import GemmType, make_2d_tma_a_desc, make_2d_tma_b_desc, make_2d_tma_d_desc, make_2d_tma_scales_a_desc
 from .tuner import jit_tuner
 from .utils import get_num_sms, ceil_div, get_col_major_tma_aligned_tensor, get_m_alignment_for_contiguous_layout
 
@@ -122,7 +122,7 @@ def get_best_configs(m: int, n: int, k: int, num_groups: int, num_sms: int,
     assert best_smem_config is not None
     assert best_num_stages is not None
 
-    # Decide the number of TMA multicast and whether broadcast on A
+    # Decide the number of TMA multicasts and whether broadcast on A
     best_tma_multicast_config = (1, True)
 
     # Try to multicast on the larger block side first
@@ -155,13 +155,13 @@ def gemm_fp8_fp8_bf16_nt(lhs: Tuple[torch.Tensor, torch.Tensor],
     Do a normal GEMM with FP8 inputs and BF16 output, with 1x128 LHS scaling and 128x128 RHS scaling.
     LHS, RHS, RHS scaling factors, and output tensors must be in contiguous format.
     RHS and RHS scaling factors are required to be transposed.
-    The LHS scaling tensor requires TMA-aligned transposed format, if your input does not match the requirement,
+    The LHS scaling tensor requires a TMA-aligned transposed format, if your input does not match the requirement,
         this function will do a transposing with a set of slow PyTorch operations.
 
     Arguments:
         lhs: the first element is an FP8 tensor (typed `torch.float8_e4m3fn`) of shape `[m, k]`,
              the second element is an FP32 1x128 scaling tensor for LHS of shape `[m, ⌈k / 128⌉]`.
-        rhs: the first element is an FP8 tensor (typed `torch.float8_e4m3fn`) of shape `[n, k]`.
+        rhs: the first element is an FP8 tensor (typed `torch.float8_e4m3fn`) of shape `[n, k]`,
              the second element is an FP32 128x128 scaling tensor for RHS of shape `[⌈n / 128⌉, ⌈k / 128⌉]`.
         out: the BF16 output tensor of shape `[m, n]`, representing the result.
     """
@@ -183,7 +183,7 @@ def gemm_fp8_fp8_bf16_nt(lhs: Tuple[torch.Tensor, torch.Tensor],
     assert out.dtype == torch.bfloat16
     assert lhs.is_contiguous() and rhs.is_contiguous() and out.is_contiguous()
 
-    # LHS scales must be transposed for TMA load, but not for RHS scales
+    # LHS scales must be transposed for TMA loads, but not for RHS scales
     # NOTES: `get_tma_aligned_lhs_scales` may launch a kernel if not processed by previous kernels
     lhs_scales = get_col_major_tma_aligned_tensor(lhs_scales)
     assert rhs_scales.is_contiguous()
@@ -201,11 +201,11 @@ def gemm_fp8_fp8_bf16_nt(lhs: Tuple[torch.Tensor, torch.Tensor],
     num_math_threads_per_group = 128
 
     tensor_map_a = make_2d_tma_a_desc(
-        GemmType.Normal, lhs, m, k, block_m, block_k)
+        GemmType.Normal, lhs, m, k, block_m, block_k, 1)
     tensor_map_b = make_2d_tma_b_desc(
-        GemmType.Normal, rhs, k, n, block_k, block_n)
+        GemmType.Normal, rhs, k, n, block_k, block_n, 1)
     tensor_map_d = make_2d_tma_d_desc(
-        GemmType.Normal, smem_config[1], out, m, n, block_m, block_n)
+        GemmType.Normal, out, m, n, block_m, block_n, 1, smem_config[1])
     tensor_map_scales_a = make_2d_tma_scales_a_desc(
         GemmType.Normal, lhs_scales, m, k, block_m, block_k)
 
@@ -237,7 +237,9 @@ def gemm_fp8_fp8_bf16_nt(lhs: Tuple[torch.Tensor, torch.Tensor],
               'NUM_TMA_MULTICAST': tma_multicast_config[0],
               'IS_TMA_MULTICAST_ON_A': tma_multicast_config[1]},
         space=(),
-        kwargs=kwargs
+        kwargs=kwargs,
+        generator=generate,
+        runtime_cls=FP8GemmRuntime,
     )
 
     # Run the kernel

deep_gemm/jit_kernels/m_grouped_gemm.py

@@ -1,9 +1,9 @@
 import torch
 from typing import Tuple
 
-from ..jit.utils import GemmType, make_2d_tma_a_desc, make_2d_tma_b_desc, make_2d_tma_d_desc, make_2d_tma_scales_a_desc
-
-from .gemm import get_best_configs, get_block_n_padding_for_smem_d
+from .gemm import get_best_configs
+from .runtime import FP8GemmRuntime, generate
+from .runtime import GemmType, make_2d_tma_a_desc, make_2d_tma_b_desc, make_2d_tma_d_desc, make_2d_tma_scales_a_desc
 from .tuner import jit_tuner
 from .utils import get_col_major_tma_aligned_tensor, get_num_sms
 
@@ -15,7 +15,7 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(lhs: Tuple[torch.Tensor, torch.Ten
     Do a grouped GEMM (contiguous format) with FP8 inputs and BF16 output, with 1x128 LHS scaling and 128x128 RHS scaling.
     LHS, RHS, RHS scaling factors, and output tensors must be in contiguous format.
     RHS and RHS scaling factors are required to be transposed.
-    The LHS scaling tensor requires TMA-aligned transposed format, if your input does not match the requirement,
+    The LHS scaling tensor requires a TMA-aligned transposed format, if your input does not match the requirement,
         this function will do a transposing with a set of slow PyTorch operations.
     On the M axis, inputs are grouped into several batches, of which batch sizes aligned to
         `get_m_alignment_for_contiguous_layout()` (128).
@@ -23,11 +23,11 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(lhs: Tuple[torch.Tensor, torch.Ten
     Arguments:
         lhs: the first element is an FP8 tensor (typed `torch.float8_e4m3fn`) of shape `[m_sum, k]`,
              the second element is an FP32 1x128 scaling tensor for LHS of shape `[m_sum, ⌈k / 128⌉]`.
-        rhs: the first element is an FP8 tensor (typed `torch.float8_e4m3fn`) of shape `[num_groups, n, k]`.
+        rhs: the first element is an FP8 tensor (typed `torch.float8_e4m3fn`) of shape `[num_groups, n, k]`,
              the second element is an FP32 128x128 scaling tensor for RHS of shape `[num_groups, ⌈n / 128⌉, ⌈k / 128⌉]`.
         out: the BF16 output tensor of shape `[m_sum, n]`, representing the result.
         m_indices: a tensor of shape `[m_sum]` with type `torch.int`.
-            `m_indices[i]` records the group which the i-th row of the LHS belong to,
+            `m_indices[i]` records the group which the i-th row of the LHS belongs to,
             which means that the i-th row of the LHS matrix will be multiplied with `rhs[m_indices[i]]`.
             Values of `m_indices` in every-m-alignment-block must also be the same.
     """
@@ -70,7 +70,7 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(lhs: Tuple[torch.Tensor, torch.Ten
     tensor_map_b = make_2d_tma_b_desc(
         GemmType.GroupedContiguous, rhs, k, n, block_k, block_n, num_groups)
     tensor_map_d = make_2d_tma_d_desc(
-        GemmType.GroupedContiguous, smem_config[1], out, m, n, block_m, block_n, num_groups)
+        GemmType.GroupedContiguous, out, m, n, block_m, block_n, num_groups, smem_config[1])
     tensor_map_scales_a = make_2d_tma_scales_a_desc(
         GemmType.GroupedContiguous, lhs_scales, m, k, block_m, block_k, num_groups)
 
@@ -103,6 +103,8 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(lhs: Tuple[torch.Tensor, torch.Ten
               'GEMM_TYPE': GemmType.GroupedContiguous},
         space=(),
         kwargs=kwargs,
+        generator=generate,
+        runtime_cls=FP8GemmRuntime,
     )
 
     # Run the kernel
@@ -116,7 +118,7 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_masked(lhs: Tuple[torch.Tensor, torch.Tensor]
     Do a grouped GEMM (masked format) with FP8 inputs and BF16 output, with 1x128 LHS scaling and 128x128 RHS scaling.
     LHS, RHS, RHS scaling factors, and output tensors must be in contiguous format.
     RHS and RHS scaling factors are required to be transposed.
-    The LHS scaling tensor requires TMA-aligned transposed format, if your input does not match the requirement,
+    The LHS scaling tensor requires a TMA-aligned transposed format, if your input does not match the requirement,
         this function will do a transposing with a set of slow PyTorch operations.
     Moreover, this alignment requirement is different with the contiguous-format kernel, as we require that each batch
         should be separately transposed.
@@ -125,7 +127,7 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_masked(lhs: Tuple[torch.Tensor, torch.Tensor]
         lhs: the first element is an FP8 tensor (typed `torch.float8_e4m3fn`) of shape `[num_groups, m_max, k]`,
              the second element is an FP32 1x128 scaling tensor for LHS of shape `[num_groups, m_max, ⌈k / 128⌉]`.
         rhs: the first element is an FP8 tensor (typed `torch.float8_e4m3fn`) of shape `[num_groups, n, k]`.
-             the second element is an FP32 128x128 scaling tensor for RHS of shape `[num_groups, ⌈n / 128⌉, ⌈k / 128⌉]`.
+             The second element is an FP32 128x128 scaling tensor for RHS of shape `[num_groups, ⌈n / 128⌉, ⌈k / 128⌉]`.
         out: the BF16 output tensor of shape `[num_groups, m_max, n]`, representing the result.
         masked_m: a tensor of shape `[num_groups]`, `masked_m[i]` records actual rows of the `lhs[i]` matrix to compute
             in the i-th group.
@@ -157,7 +159,6 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_masked(lhs: Tuple[torch.Tensor, torch.Tensor]
     assert rhs_scales.is_contiguous()
 
     # Auto-tuning with compilation
-    global includes, template
     num_sms = get_num_sms()
     num_sms, block_m, block_n, num_stages, tma_multicast_config, smem_config = get_best_configs(
         expected_m, n, k, num_groups, num_sms, is_grouped_masked=True)
@@ -175,7 +176,7 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_masked(lhs: Tuple[torch.Tensor, torch.Tensor]
     tensor_map_b = make_2d_tma_b_desc(
         GemmType.GroupedMasked, rhs, k, n, block_k, block_n, num_groups)
     tensor_map_d = make_2d_tma_d_desc(
-        GemmType.GroupedMasked, smem_config[1], out, m, n, block_m, block_n, num_groups)
+        GemmType.GroupedMasked, out, m, n, block_m, block_n, num_groups, smem_config[1])
     tensor_map_scales_a = make_2d_tma_scales_a_desc(
         GemmType.GroupedMasked, lhs_scales, m, k, block_m, block_k, num_groups)
 
@@ -208,6 +209,8 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_masked(lhs: Tuple[torch.Tensor, torch.Tensor]
               'GEMM_TYPE': GemmType.GroupedMasked},
         space=(),
         kwargs=kwargs,
+        generator=generate,
+        runtime_cls=FP8GemmRuntime,
     )
 
     # Run the kernel

deep_gemm/jit_kernels/runtime.py

@@ -0,0 +1,256 @@
+import ctypes
+import os
+import enum
+import torch
+import cuda.bindings.driver as cbd
+from typing import Any, Dict, Tuple
+
+from ..jit.runtime import Runtime
+
+
+def generate(**kwargs: Dict[str, Any]) -> str:
+    code = f'''
+#ifdef __CUDACC_RTC__
+#include <deep_gemm/nvrtc_std.cuh>
+#else
+#include <cuda.h>
+#include <string>
+#endif
+
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+
+#include <deep_gemm/fp8_gemm.cuh>
+
+using namespace deep_gemm;
+
+__global__ void dummy_kernel() {{
+  void *ptr = (void *)&fp8_gemm_kernel<
+    {kwargs['N']},
+    {kwargs['K']},
+    {kwargs['BLOCK_M']},
+    {kwargs['BLOCK_N']},
+    {kwargs['BLOCK_K']},
+    {kwargs['BLOCK_N_PADDING']},
+    {kwargs['SWIZZLE_D_MODE']},
+    {kwargs['NUM_GROUPS']},
+    {kwargs['NUM_STAGES']},
+    {kwargs['NUM_TMA_THREADS']},
+    {kwargs['NUM_MATH_THREADS_PER_GROUP']},
+    {kwargs['NUM_TMA_MULTICAST']},
+    {'true' if kwargs['IS_TMA_MULTICAST_ON_A'] else 'false'},
+    GemmType::{kwargs['GEMM_TYPE']}
+  >;
+}}
+'''
+
+    if int(os.getenv('DG_JIT_DEBUG', 0)):
+        print(f'Generated code:\n{code}')
+    return code
+
+
+class FP8GemmRuntime(Runtime):
+    def __init__(self, path: str) -> None:
+        super().__init__(path, 'fp8_gemm', launch, [
+            'NUM_TMA_MULTICAST',
+            'M',
+            'BLOCK_M',
+            'GMEM_D',
+            'SCALES_B',
+            'GROUPED_LAYOUT',
+            'NUM_SMS',
+            'SMEM_SIZE',
+            'TENSOR_MAP_A',
+            'TENSOR_MAP_B',
+            'TENSOR_MAP_SCALES_A',
+            'TENSOR_MAP_D',
+            'STREAM',
+        ])
+
+
+class Layout(enum.Enum):
+    RowMajor = 0
+    ColMajor = 1
+
+
+class GemmType(enum.Enum):
+    Normal = 0
+    GroupedContiguous = 1
+    GroupedMasked = 2
+
+    def __str__(self) -> str:
+        return {
+            0: 'Normal',
+            1: 'GroupedContiguous',
+            2: 'GroupedMasked',
+        }[self.value]
+
+
+tmap_type_map: Dict[Any, str] = {
+    torch.int8:            cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8,
+    torch.int16:           cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT16,
+    torch.int32:           cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_INT32,
+    torch.int64:           cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_INT64,
+    torch.uint8:           cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8,
+    torch.uint16:          cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT16,
+    torch.uint32:          cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT32,
+    torch.uint64:          cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT64,
+    torch.float32:         cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_FLOAT32,
+    torch.float16:         cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_FLOAT16,
+    torch.bfloat16:        cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16,
+    torch.float8_e4m3fn:   cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8,
+    torch.float8_e4m3fnuz: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8,
+    torch.float8_e5m2:     cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8,
+    torch.float8_e5m2fnuz: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8,
+}
+
+swizzle_type_map = {
+    0:   cbd.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_NONE,
+    32:  cbd.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_32B,
+    64:  cbd.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_64B,
+    128: cbd.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B,
+}
+
+
+def get_num_math_warpgroups(block_m: int) -> int:
+    return 1 if block_m == 64 else 2
+
+
+def get_num_threads_per_sm(num_tma_threads: int, num_math_threads_per_group: int, block_m: int) -> int:
+    assert num_math_threads_per_group == 128, 'Only support 128 threads per math group'
+    return get_num_math_warpgroups(block_m) * num_math_threads_per_group + num_tma_threads
+
+
+def make_2d_tma_copy_desc(global_address: torch.Tensor,
+                          gmem_dim: Tuple[cbd.cuuint64_t, cbd.cuuint64_t],
+                          stride_in_bytes: cbd.cuuint64_t,
+                          smem_dim: Tuple[cbd.cuuint32_t, cbd.cuuint32_t],
+                          swizzle_type: cbd.CUtensorMapSwizzle) -> cbd.CUtensorMap:
+    tensor_dtype = tmap_type_map[global_address.dtype]
+    res, tensor_map = cbd.cuTensorMapEncodeTiled(
+        tensor_dtype,
+        2,
+        global_address.data_ptr(),
+        gmem_dim,
+        (stride_in_bytes, ),
+        smem_dim,
+        (cbd.cuuint32_t(1), cbd.cuuint32_t(1)),
+        cbd.CUtensorMapInterleave.CU_TENSOR_MAP_INTERLEAVE_NONE,
+        swizzle_type,
+        cbd.CUtensorMapL2promotion.CU_TENSOR_MAP_L2_PROMOTION_L2_256B,
+        cbd.CUtensorMapFloatOOBfill.CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE,
+    )
+
+    if res != cbd.CUresult.CUDA_SUCCESS:
+        raise Exception(f'Failed to encode tensor map: {res}')
+    return tensor_map
+
+
+def make_2d_tma_desc(global_address: torch.Tensor, layout: Layout,
+                     gmem_rows: int, gmem_cols: int,
+                     smem_rows: int, smem_cols: int,
+                     swizzle_type: cbd.CUtensorMapSwizzle = cbd.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B) -> cbd.CUtensorMap:
+    if layout == Layout.RowMajor:
+        gmem_dim = (cbd.cuuint64_t(gmem_cols), cbd.cuuint64_t(gmem_rows))
+        smem_dim = (cbd.cuuint32_t(smem_cols), cbd.cuuint32_t(smem_rows))
+        return make_2d_tma_copy_desc(global_address, gmem_dim, cbd.cuuint64_t(gmem_cols * global_address.element_size()), smem_dim, swizzle_type)
+    else:
+        gmem_dim = (cbd.cuuint64_t(gmem_rows), cbd.cuuint64_t(gmem_cols))
+        smem_dim = (cbd.cuuint32_t(smem_rows), cbd.cuuint32_t(smem_cols))
+        return make_2d_tma_copy_desc(global_address, gmem_dim, cbd.cuuint64_t(gmem_rows * global_address.element_size()), smem_dim, swizzle_type)
+
+
+def make_2d_tma_a_desc(gemm_type: GemmType, global_address: torch.Tensor,
+                       shape_m: int, shape_k: int,
+                       block_m: int, block_k: int,
+                       num_groups: int) -> cbd.CUtensorMap:
+    return make_2d_tma_desc(global_address, Layout.RowMajor,
+                            shape_m * (num_groups if gemm_type == GemmType.GroupedMasked else 1), shape_k,
+                            block_m, block_k)
+
+
+def make_2d_tma_b_desc(gemm_type: GemmType, global_address: torch.Tensor,
+                       shape_k: int, shape_n: int,
+                       block_k: int, block_n: int,
+                       num_groups: int) -> cbd.CUtensorMap:
+    return make_2d_tma_desc(global_address, Layout.ColMajor,
+                            shape_k, shape_n * (num_groups if gemm_type != GemmType.Normal else 1),
+                            block_k, block_n)
+
+
+def make_2d_tma_d_desc(gemm_type: GemmType, global_address: torch.Tensor,
+                       shape_m: int, shape_n: int,
+                       block_m: int, block_n: int,
+                       num_groups: int, swizzle_mode: int) -> cbd.CUtensorMap:
+    # Swizzling requires the inner box dim to be less or equal than `kSwizzleDMode`
+    # bytes, so `BLOCK_N * sizeof(T) / kSwizzleDMode` TMA stores are required
+    return make_2d_tma_desc(global_address, Layout.RowMajor,
+                            shape_m * (num_groups if gemm_type == GemmType.GroupedMasked else 1), shape_n,
+                            block_m, block_n if swizzle_mode == 0 else swizzle_mode // global_address.element_size(),
+                            swizzle_type_map[swizzle_mode])
+
+
+def make_2d_tma_scales_a_desc(gemm_type: GemmType, global_address: torch.Tensor, shape_m: int, shape_k: int, block_m: int, block_k: int, num_groups: int = 1) -> cbd.CUtensorMap:
+    # Make TMA aligned to 16 bytes
+    tma_alignment = 16 / global_address.element_size()
+    shape_m = (shape_m + tma_alignment - 1) // tma_alignment * tma_alignment
+
+    return make_2d_tma_desc(global_address, Layout.ColMajor,
+                            shape_m, (shape_k + block_k - 1) // block_k * (num_groups if gemm_type == GemmType.GroupedMasked else 1),
+                            block_m, 1, cbd.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_NONE)
+
+
+def launch(kernel: cbd.CUkernel, num_tma_multicast: int, shape_m: int,
+           block_m: int, gmem_d: torch.Tensor, scales_b: torch.Tensor,
+           grouped_layout: torch.Tensor, num_sms: int, smem_size: int,
+           tensor_map_a: cbd.CUtensorMap, tensor_map_b: cbd.CUtensorMap,
+           tensor_map_scales_a: cbd.CUtensorMap, tensor_map_d: cbd.CUtensorMap,
+           stream: cbd.CUstream) -> cbd.CUresult:
+    num_tma_threads = 128
+    num_math_threads_per_group = 128
+
+    res = cbd.cuKernelSetAttribute(cbd.CUfunction_attribute.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem_size, kernel, cbd.CUdevice(gmem_d.device.index))[0]
+    if res != cbd.CUresult.CUDA_SUCCESS:
+        raise Exception(f'Failed to set max dynamic shared memory size: {res}')
+
+    attr_val = cbd.CUlaunchAttributeValue()
+    attr_val.clusterDim.x = num_tma_multicast
+    attr_val.clusterDim.y = 1
+    attr_val.clusterDim.z = 1
+    attr = cbd.CUlaunchAttribute()
+    attr.id = cbd.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION
+    attr.value = attr_val
+
+    config = cbd.CUlaunchConfig()
+    config.numAttrs = 1
+    config.attrs = [attr]
+    config.gridDimX = num_sms
+    config.gridDimY = 1
+    config.gridDimZ = 1
+    config.blockDimX = get_num_threads_per_sm(num_tma_threads, num_math_threads_per_group, block_m)
+    config.blockDimY = 1
+    config.blockDimZ = 1
+    config.sharedMemBytes = smem_size
+    config.hStream = stream
+
+    arg_values = (
+        gmem_d.data_ptr(),
+        scales_b.data_ptr(),
+        grouped_layout.data_ptr(),
+        shape_m,
+        tensor_map_a,
+        tensor_map_b,
+        tensor_map_scales_a,
+        tensor_map_d,
+    )
+    arg_types = (
+        ctypes.c_void_p,
+        ctypes.c_void_p,
+        ctypes.c_void_p,
+        ctypes.c_uint32,
+        None,
+        None,
+        None,
+        None,
+    )
+    return cbd.cuLaunchKernelEx(config, kernel, (arg_values, arg_types), 0)

deep_gemm/jit_kernels/tuner.py

@@ -1,29 +1,28 @@
 import copy
 import os
 import torch
-from typing import Any, Dict
+import cuda.bindings.driver as cbd
+from typing import Any, Callable, Dict, Type, Tuple
 
-import cuda.bindings.driver as cuda
-
-from ..jit import build, generate, Runtime
+from ..jit import build, Runtime
 
 
 class JITTuner:
     def __init__(self) -> None:
         self.tuned = {}
 
-    def compile_and_tune(self, name: str, keys: Dict[str, Any], space: tuple, kwargs: Dict[str, Any]) -> Runtime:
-        # NOTES: we always assume the space and template will not change
-        # We also assume the GPU device will not be changed
+    def compile_and_tune(self, name: str, keys: Dict[str, Any], space: tuple,
+                         kwargs: Dict[str, Any], generator: Callable[..., str], runtime_cls: Type[Runtime]) -> Tuple[Runtime, Dict[str, Any]]:
+        # NOTES: we always assume the space, template and GPU devices will not change
         # NOTES: the function must have no accumulated side effects
         keys = {k: keys[k] for k in sorted(keys.keys())}
         signature = (name, f'{keys}')
         if signature in self.tuned:
-            if os.getenv('DG_JIT_DEBUG', None):
+            if int(os.getenv('DG_JIT_DEBUG', 0)):
                 print(f'Using cached JIT kernel {name} with keys {keys}')
             return self.tuned[signature]
 
-        if os.getenv('DG_JIT_DEBUG', None):
+        if int(os.getenv('DG_JIT_DEBUG', 0)):
             print(f'Auto-tuning JIT kernel {name} with keys {keys}')
 
         assert signature not in self.tuned
@@ -35,19 +34,19 @@ class JITTuner:
             assert isinstance(tuned_keys, dict)
             full_keys = copy.deepcopy(keys)
             full_keys.update(tuned_keys)
-            code = generate(**kwargs, **full_keys)
-            kernels.append((build(name, code), full_keys))
+            code = generator(**kwargs, **full_keys)
+            kernels.append((build(name, code, runtime_cls), full_keys))
 
+        # TODO: fix tuning with space > 1
         best_runtime, best_time, best_keys = None, None, None
         for runtime, tuned_keys in kernels:
             if len(space) > 1:
                 # Check kernel validity
                 return_code = runtime(**tuned_keys, **kwargs)
-                if return_code != cuda.CUresult.CUDA_SUCCESS:
-                    # Pass illegal kernels, e.g. insufficient shared memory capacity
-                    if os.getenv('DG_JIT_DEBUG', None):
-                        print(
-                            f'Illegal JIT kernel {name} with keys {keys} and tuned keys {tuned_keys}: error code {return_code}')
+                if return_code != cbd.CUresult.CUDA_SUCCESS:
+                    # Pass illegal kernels, e.g., insufficient shared memory capacity
+                    if int(os.getenv('DG_JIT_DEBUG', 0)):
+                        print(f'Illegal JIT kernel {name} with keys {keys} and tuned keys {tuned_keys}: error code {return_code}')
                     continue
 
                 # Measure performance with L2 flush and a large GEMM kernel before to reduce overhead between kernels
@@ -59,7 +58,7 @@ class JITTuner:
                     (8192, 8192), dtype=torch.float, device='cuda')
                 start_event.record()
                 for i in range(20):
-                    assert runtime(**tuned_keys, **kwargs) == cuda.CUresult.CUDA_SUCCESS
+                    assert runtime(**tuned_keys, **kwargs) == cbd.CUresult.CUDA_SUCCESS
                 end_event.record()
                 end_event.synchronize()
                 elapsed_time = start_event.elapsed_time(end_event)
@@ -69,13 +68,12 @@ class JITTuner:
             # Compare if better
             if best_time is None or elapsed_time < best_time:
                 best_runtime, best_time, best_keys = runtime, elapsed_time, tuned_keys
-            if os.getenv('DG_JIT_DEBUG', None):
-                print(
-                    f'Tuned JIT kernel {name} with keys {keys} and tuned keys {tuned_keys} has time {elapsed_time}')
+            if int(os.getenv('DG_JIT_DEBUG', 0)):
+                print(f'Tuned JIT kernel {name} with keys {keys} and tuned keys {tuned_keys} has time {elapsed_time}')
         assert best_runtime is not None, f'Failed to tune JIT kernel {name} with keys {keys}'
 
         # Cache the best runtime and return
-        if os.getenv('DG_JIT_DEBUG', None) or os.getenv('DG_PRINT_AUTOTUNE', None):
+        if int(os.getenv('DG_JIT_DEBUG', 0)) or int(os.getenv('DG_PRINT_AUTOTUNE', 0)):
             print(
                 f'Best JIT kernel {name} with keys {keys} has tuned keys {best_keys} and time {best_time}')
         self.tuned[signature] = (best_runtime, best_keys)