leavelet/deepgemm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: support libtorch

Browse Source
This commit is contained in:
FlamingoPg
2025-10-21 09:55:33 +00:00
parent 239112cb4c
commit 301cbc1d75
22 changed files with 652 additions and 245 deletions
Download Patch File Download Diff File Expand all files Collapse all files
deep_gemm/utils/layout.py
+12 -10
View File
@@ -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