deep_gemm/jit_kernels/__init__.py

from .gemm import gemm_fp8_fp8_bf16_nt
from .m_grouped_gemm import (
    m_grouped_gemm_fp8_fp8_bf16_nt_contiguous,
    m_grouped_gemm_fp8_fp8_bf16_nt_masked
)
from .wgrad_gemm import (
    wgrad_gemm_fp8_fp8_fp32_nt,
    k_grouped_wgrad_gemm_fp8_fp8_fp32_nt
)
from .utils import (
    ceil_div, set_num_sms, get_num_sms,
    get_col_major_tma_aligned_tensor,
    get_m_alignment_for_contiguous_layout
)
Initial commit 2025-02-25 22:52:41 +08:00			`from .gemm import gemm_fp8_fp8_bf16_nt`
			`from .m_grouped_gemm import (`
			`m_grouped_gemm_fp8_fp8_bf16_nt_contiguous,`
			`m_grouped_gemm_fp8_fp8_bf16_nt_masked`
			`)`
Weight gradient kernels for dense and MoE models (#95) * Init weight gradient kernels. * Support unaligned n,k and gmem stride * Update docs * Several cleanups * Remove restrictions on N * Add stride(0) assertions --------- Co-authored-by: Chenggang Zhao <chenggangz@deepseek.com> 2025-05-14 14:47:58 +08:00			`from .wgrad_gemm import (`
			`wgrad_gemm_fp8_fp8_fp32_nt,`
			`k_grouped_wgrad_gemm_fp8_fp8_fp32_nt`
			`)`
Initial commit 2025-02-25 22:52:41 +08:00			`from .utils import (`
fix typo 2025-02-26 18:37:22 +08:00			`ceil_div, set_num_sms, get_num_sms,`
Initial commit 2025-02-25 22:52:41 +08:00			`get_col_major_tma_aligned_tensor,`
			`get_m_alignment_for_contiguous_layout`
			`)`