feat: support libtorch

2025-08-26 19:17:08 -07:00
parent 239112cb4c
commit 301cbc1d75
22 changed files with 652 additions and 245 deletions
--- a/deep_gemm/init.py
+++ b/deep_gemm/init.py
@@ -1,5 +1,7 @@
 import os
 import subprocess
+import torch
+import torch.utils.cpp_extension

 # Set some default environment provided at setup
 try:
@@ -11,62 +13,12 @@ try:
 except ImportError:
    pass

-# Configs
-import deep_gemm_cpp
-from deep_gemm_cpp import (
-    set_num_sms,
-    get_num_sms,
-    set_tc_util,
-    get_tc_util,
-)
-
-# Kernels
-from deep_gemm_cpp import (
-    # FP8 GEMMs
-    fp8_gemm_nt, fp8_gemm_nn,
-    fp8_gemm_tn, fp8_gemm_tt,
-    fp8_gemm_nt_skip_head_mid,
-    m_grouped_fp8_gemm_nt_contiguous,
-    m_grouped_fp8_gemm_nn_contiguous,
-    m_grouped_fp8_gemm_nt_masked,
-    k_grouped_fp8_gemm_nt_contiguous,
-    k_grouped_fp8_gemm_tn_contiguous,
-    # BF16 GEMMs
-    bf16_gemm_nt, bf16_gemm_nn,
-    bf16_gemm_tn, bf16_gemm_tt,
-    m_grouped_bf16_gemm_nt_contiguous,
-    m_grouped_bf16_gemm_nt_masked,
-    # cuBLASLt GEMMs
-    cublaslt_gemm_nt, cublaslt_gemm_nn,
-    cublaslt_gemm_tn, cublaslt_gemm_tt,
-    # Einsum kernels
-    einsum,
-    # Attention kernels
-    fp8_mqa_logits,
-    get_paged_mqa_logits_metadata,
-    fp8_paged_mqa_logits,
-    # Layout kernels
-    transform_sf_into_required_layout
-)
-
-# Some alias for legacy supports
-# TODO: remove these later
-fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_gemm_nt_masked
-bf16_m_grouped_gemm_nt_masked = m_grouped_bf16_gemm_nt_masked
-
-# Some utils
-from . import testing
-from . import utils
-from .utils import *
+from . import deep_gemm_cpp  # noqa: F401  # Registers ops into torch.ops without touching CUDA


-# Initialize CPP modules
 def _find_cuda_home() -> str:
-    # TODO: reuse PyTorch API later
-    # For some PyTorch versions, the original `_find_cuda_home` will initialize CUDA, which is incompatible with process forks
    cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
    if cuda_home is None:
-        # noinspection PyBroadException
        try:
            with open(os.devnull, 'w') as devnull:
                nvcc = subprocess.check_output(['which', 'nvcc'], stderr=devnull).decode().rstrip('\r\n')
@@ -79,7 +31,106 @@ def _find_cuda_home() -> str:
    return cuda_home


-deep_gemm_cpp.init(
-    os.path.dirname(os.path.abspath(__file__)), # Library root directory path
-    _find_cuda_home()                           # CUDA home
-)
+# Lazy runtime init to be fork-safe on Linux (avoid initializing CUDA before fork)
+_dg_initialized = False
+
+def _ensure_initialized() -> None:
+    global _dg_initialized
+    if _dg_initialized:
+        return
+    library_root = os.path.dirname(os.path.abspath(__file__))
+    torch.ops.deep_gemm.init(library_root, _find_cuda_home())
+    _dg_initialized = True
+
+
+def _wrap_op(name: str):
+    def _fn(*args, **kwargs):
+        _ensure_initialized()
+        return getattr(torch.ops.deep_gemm, name)(*args, **kwargs)
+    return _fn
+
+set_num_sms = _wrap_op('set_num_sms')
+get_num_sms = _wrap_op('get_num_sms')
+set_tc_util = _wrap_op('set_tc_util')
+get_tc_util = _wrap_op('get_tc_util')
+
+fp8_gemm_nt = _wrap_op('fp8_gemm_nt')
+fp8_gemm_nn = _wrap_op('fp8_gemm_nn')
+fp8_gemm_tn = _wrap_op('fp8_gemm_tn')
+fp8_gemm_tt = _wrap_op('fp8_gemm_tt')
+m_grouped_fp8_gemm_nt_contiguous = _wrap_op('m_grouped_fp8_gemm_nt_contiguous')
+m_grouped_fp8_gemm_nn_contiguous = _wrap_op('m_grouped_fp8_gemm_nn_contiguous')
+# Export both canonical name and backward-compat alias
+m_grouped_fp8_gemm_nt_masked = _wrap_op('m_grouped_fp8_gemm_nt_masked')
+fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_gemm_nt_masked
+k_grouped_fp8_gemm_nt_contiguous = _wrap_op('k_grouped_fp8_gemm_nt_contiguous')
+k_grouped_fp8_gemm_tn_contiguous = _wrap_op('k_grouped_fp8_gemm_tn_contiguous')
+
+# BF16 GEMMs
+bf16_gemm_nt = _wrap_op('bf16_gemm_nt')
+bf16_gemm_nn = _wrap_op('bf16_gemm_nn')
+bf16_gemm_tn = _wrap_op('bf16_gemm_tn')
+bf16_gemm_tt = _wrap_op('bf16_gemm_tt')
+m_grouped_bf16_gemm_nt_contiguous = _wrap_op('m_grouped_bf16_gemm_nt_contiguous')
+m_grouped_bf16_gemm_nt_masked = _wrap_op('m_grouped_bf16_gemm_nt_masked')
+
+# cuBLASLt GEMMs
+cublaslt_gemm_nt = _wrap_op('cublaslt_gemm_nt')
+cublaslt_gemm_nn = _wrap_op('cublaslt_gemm_nn')
+cublaslt_gemm_tn = _wrap_op('cublaslt_gemm_tn')
+cublaslt_gemm_tt = _wrap_op('cublaslt_gemm_tt')
+
+# Attention kernel
+fp8_gemm_nt_skip_head_mid = _wrap_op('fp8_gemm_nt_skip_head_mid')
+fp8_mqa_logits = _wrap_op('fp8_mqa_logits')
+get_paged_mqa_logits_metadata = _wrap_op('get_paged_mqa_logits_metadata')
+fp8_paged_mqa_logits = _wrap_op('fp8_paged_mqa_logits')
+
+# Einsum kernel
+einsum = _wrap_op('einsum')
+
+# Layout kernels
+transform_sf_into_required_layout = _wrap_op('transform_sf_into_required_layout')
+
+# Utility functions
+get_tma_aligned_size = _wrap_op('get_tma_aligned_size')
+get_mk_alignment_for_contiguous_layout = _wrap_op('get_mk_alignment_for_contiguous_layout')
+get_mn_major_tma_aligned_tensor = _wrap_op('get_mn_major_tma_aligned_tensor')
+get_mn_major_tma_aligned_packed_ue8m0_tensor = _wrap_op('get_mn_major_tma_aligned_packed_ue8m0_tensor')
+get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor = _wrap_op('get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor')
+
+# Some utils
+from . import testing
+from . import utils
+from .utils import *
+
+def _verify_ops_loaded():
+    expected_ops = [
+        'init', 'set_num_sms', 'get_num_sms', 'set_tc_util', 'get_tc_util',
+        'fp8_gemm_nt', 'fp8_gemm_nn', 'fp8_gemm_tn', 'fp8_gemm_tt',
+        'm_grouped_fp8_gemm_nt_contiguous', 'm_grouped_fp8_gemm_nn_contiguous',
+        'm_grouped_fp8_gemm_nt_masked', 'k_grouped_fp8_gemm_nt_contiguous',
+        'k_grouped_fp8_gemm_tn_contiguous',
+        'transform_sf_into_required_layout', 'get_tma_aligned_size',
+        'get_mk_alignment_for_contiguous_layout', 'get_mn_major_tma_aligned_tensor',
+        'get_mn_major_tma_aligned_packed_ue8m0_tensor',
+        'get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor',
+        'fp8_gemm_nt_skip_head_mid', 'fp8_mqa_logits',
+        'get_paged_mqa_logits_metadata', 'fp8_paged_mqa_logits',
+        'einsum',
+        'cublaslt_gemm_nt', 'cublaslt_gemm_nn',
+        'cublaslt_gemm_tn', 'cublaslt_gemm_tt',
+    ]
+    
+    available_ops = list(torch.ops.deep_gemm.__dict__.keys())
+    missing_ops = [op for op in expected_ops if op not in available_ops]
+    
+    if missing_ops:
+        print(f"Warning: Missing operations: {missing_ops}")
+
+
+_ensure_initialized()
+
+
+if __debug__:
+    _verify_ops_loaded()
--- a/deep_gemm/utils/layout.py
+++ b/deep_gemm/utils/layout.py
@@ -1,11 +1,13 @@
-from deep_gemm_cpp import (
-    get_tma_aligned_size,
-    get_mk_alignment_for_contiguous_layout,
-    get_mn_major_tma_aligned_tensor,
-    get_mn_major_tma_aligned_packed_ue8m0_tensor,
-    get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor
-)
+import torch
+from .. import _ensure_initialized

-# Some alias
-get_m_alignment_for_contiguous_layout = get_mk_alignment_for_contiguous_layout
-get_k_alignment_for_contiguous_layout = get_mk_alignment_for_contiguous_layout
+_ensure_initialized()
+
+get_tma_aligned_size = torch.ops.deep_gemm.get_tma_aligned_size
+get_mk_alignment_for_contiguous_layout = torch.ops.deep_gemm.get_mk_alignment_for_contiguous_layout
+get_mn_major_tma_aligned_tensor = torch.ops.deep_gemm.get_mn_major_tma_aligned_tensor
+get_mn_major_tma_aligned_packed_ue8m0_tensor = torch.ops.deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor
+get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor = torch.ops.deep_gemm.get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor
+
+get_m_alignment_for_contiguous_layout = torch.ops.deep_gemm.get_mk_alignment_for_contiguous_layout
+get_k_alignment_for_contiguous_layout = torch.ops.deep_gemm.get_mk_alignment_for_contiguous_layout