DeepSeek/DeepGEMM
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepGEMM synced 2025-06-26 23:15:49 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add DG_PRINT_CONFIGS

Browse Source
This commit is contained in:
Chenggang Zhao
2025-05-15 16:36:40 +08:00
parent 816b39053a
commit 4373af2e82
6 changed files with 26 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
deep_gemm/jit_kernels/gemm.py
View File

@@ -238,5 +238,5 @@ def gemm_fp8_fp8_bf16_nt(lhs: Tuple[torch.Tensor, torch.Tensor],
# Generate, build and run the kernel
code = FP8GemmRuntime.generate(**kwargs)
runtime = build('gemm_fp8_fp8_bf16_nt', code, FP8GemmRuntime)
runtime = build('gemm_fp8_fp8_bf16_nt', code, FP8GemmRuntime, kwargs)
runtime(**kwargs)

4
deep_gemm/jit_kernels/m_grouped_gemm.py
View File

@@ -102,7 +102,7 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(lhs: Tuple[torch.Tensor, torch.Ten
# Generate, build and run the kernel
code = FP8GemmRuntime.generate(**kwargs)
runtime = build('m_grouped_gemm_fp8_fp8_bf16_nt', code, FP8GemmRuntime)
runtime = build('m_grouped_gemm_fp8_fp8_bf16_nt', code, FP8GemmRuntime, kwargs)
runtime(**kwargs)
@@ -201,5 +201,5 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_masked(lhs: Tuple[torch.Tensor, torch.Tensor]
# Generate, build and run the kernel
code = FP8GemmRuntime.generate(**kwargs)
runtime = build('m_grouped_gemm_fp8_fp8_bf16_nt', code, FP8GemmRuntime)
runtime = build('m_grouped_gemm_fp8_fp8_bf16_nt', code, FP8GemmRuntime, kwargs)
runtime(**kwargs)

2
deep_gemm/jit_kernels/wgrad_gemm.py
View File

@@ -111,7 +111,7 @@ def wgrad_gemm_fp8_fp8_fp32_nt(lhs: Tuple[torch.Tensor, torch.Tensor],
# Generate, build and run the kernel
code = FP8WGradGemmRuntime.generate(**kwargs)
runtime = build('wgrad_gemm_fp8_fp8_fp32_nt', code, FP8WGradGemmRuntime)
runtime = build('wgrad_gemm_fp8_fp8_fp32_nt', code, FP8WGradGemmRuntime, kwargs)
runtime(**kwargs)