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_AUTOTUNE to README

Browse Source
This commit is contained in:
Chenggang Zhao 2025-05-07 11:46:52 +08:00
parent daec8fd2fc
commit 085b4a1532
2 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
README.md
View File

@ -120,6 +120,8 @@ The library also provides some environment variables, which may be useful:
- `DG_JIT_PRINT_COMPILER_COMMAND`: `0` or `1`, print NVCC compilation command, `0` by default - `DG_JIT_PRINT_COMPILER_COMMAND`: `0` or `1`, print NVCC compilation command, `0` by default
- Post optimization - Post optimization
- `DG_JIT_DISABLE_FFMA_INTERLEAVE`: `0` or `1`, disable FFMA-interleaving optimization, `0` by default - `DG_JIT_DISABLE_FFMA_INTERLEAVE`: `0` or `1`, disable FFMA-interleaving optimization, `0` by default
- Heuristic selection
- `DG_PRINT_AUTOTUNE`: `0` or `1`, print selected configs for each shape, `0` by default
- Testing - Testing
- `DG_NSYS_PROFILING`: `0` or `1`, Nsight-system compatible testing, `0` by default - `DG_NSYS_PROFILING`: `0` or `1`, Nsight-system compatible testing, `0` by default

3
deep_gemm/jit_kernels/tuner.py
View File

@ -74,8 +74,7 @@ class JITTuner:
# Cache the best runtime and return # Cache the best runtime and return
if int(os.getenv('DG_JIT_DEBUG', 0)) or int(os.getenv('DG_PRINT_AUTOTUNE', 0)): if int(os.getenv('DG_JIT_DEBUG', 0)) or int(os.getenv('DG_PRINT_AUTOTUNE', 0)):
print( print(f'Best JIT kernel {name} with keys {keys} has tuned keys {best_keys} and time {best_time}')
f'Best JIT kernel {name} with keys {keys} has tuned keys {best_keys} and time {best_time}')
self.tuned[signature] = (best_runtime, best_keys) self.tuned[signature] = (best_runtime, best_keys)
return best_runtime, best_keys return best_runtime, best_keys