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

Fix JIT tests

Browse Source
This commit is contained in:
Chenggang Zhao
2025-05-16 14:39:58 +08:00
parent 78d8362e7a
commit 391755ada0
6 changed files with 14 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
deep_gemm/jit/compiler.py
View File

@@ -166,7 +166,7 @@ class Compiler:
os.replace(tmp_cubin_path, cubin_path) os.replace(tmp_cubin_path, cubin_path)
# Put cache and return # Put cache and return
runtime = runtime_cache.get(path, runtime_cls, name, kwargs) runtime = runtime_cache.get(path, runtime_cls, name, kwargs, force_enable_cache=True)
assert runtime is not None assert runtime is not None
return runtime return runtime

10
deep_gemm/jit/runtime.py
View File

@@ -33,7 +33,7 @@ class Runtime:
def launch(kernel: cbd.CUkernel, kwargs: Dict[str, Any]) -> cbd.CUresult: def launch(kernel: cbd.CUkernel, kwargs: Dict[str, Any]) -> cbd.CUresult:
raise NotImplemented raise NotImplemented
def __call__(self, kwargs: Dict[str, Any]) -> cbd.CUresult: def __call__(self, **kwargs) -> cbd.CUresult:
# Load CUBIN # Load CUBIN
if self.kernel is None: if self.kernel is None:
start_time = time.time_ns() start_time = time.time_ns()
@@ -81,17 +81,19 @@ class RuntimeCache:
self.cache[path] = runtime self.cache[path] = runtime
def get(self, path: str, runtime_cls: Type[Runtime], def get(self, path: str, runtime_cls: Type[Runtime],
name: str = '', kwargs: Dict[str, Any] = None) -> Optional[Runtime]: name: str = '', kwargs: Dict[str, Any] = None,
force_enable_cache: bool = False) -> Optional[Runtime]:
# In Python runtime # In Python runtime
if path in self.cache: if path in self.cache:
return self.cache[path] return self.cache[path]
# Already compiled # Already compiled
if not int(os.getenv('DG_JIT_DISABLE_CACHE', 0)) and os.path.exists(path) and Runtime.is_path_valid(path): use_cache = force_enable_cache or not int(os.getenv('DG_JIT_DISABLE_CACHE', 0))
if use_cache and os.path.exists(path) and Runtime.is_path_valid(path):
# Print heuristic for the first time # Print heuristic for the first time
if name and (int(os.getenv('DG_JIT_DEBUG', 0)) or int(os.getenv('DG_PRINT_CONFIGS', 0))): if name and (int(os.getenv('DG_JIT_DEBUG', 0)) or int(os.getenv('DG_PRINT_CONFIGS', 0))):
simplified_kwargs = dict() simplified_kwargs = dict()
for key, value in kwargs.items(): for key, value in kwargs.items() if kwargs is not None else dict().items():
value = f'torch.Tensor<{value.dtype}>' if isinstance(value, torch.Tensor) else value value = f'torch.Tensor<{value.dtype}>' if isinstance(value, torch.Tensor) else value
value = f'cuda.bindings.driver.CUtensorMap' if isinstance(value, cbd.CUtensorMap) else value value = f'cuda.bindings.driver.CUtensorMap' if isinstance(value, cbd.CUtensorMap) else value
simplified_kwargs[key] = value simplified_kwargs[key] = value

2
deep_gemm/jit_kernels/gemm.py
View File

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

4
deep_gemm/jit_kernels/m_grouped_gemm.py
View File

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

2
deep_gemm/jit_kernels/wgrad_gemm.py
View File

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

5
tests/test_jit.py
View File

@@ -74,7 +74,8 @@ static void __instantiate_kernel() {{
if __name__ == '__main__': if __name__ == '__main__':
print('Generated code:') print('Generated code:')
code = VectorAddRuntime.generate(T='float') kwargs = {'T': 'float'}
code = VectorAddRuntime.generate(kwargs)
print(code) print(code)
print() print()
@@ -85,7 +86,7 @@ if __name__ == '__main__':
# Build # Build
print('Building ...') print('Building ...')
func = compiler_cls.build('test_func', code, VectorAddRuntime) func = compiler_cls.build('test_func', code, VectorAddRuntime, kwargs)
# Run and check # Run and check
a = torch.randn((1024, ), dtype=torch.float32, device='cuda') a = torch.randn((1024, ), dtype=torch.float32, device='cuda')