ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Make autotune default

Browse source
This commit is contained in:
Datta Nimmaturi 2026-01-13 05:40:55 +00:00
commit 1d86812b7c
1 changed files with 3 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

13
unsloth/kernels/moe/autotune_cache.py
View file

@ -147,17 +147,10 @@ def get_or_autotune_moe_kernels(
seq_len,
)
# 0. Check for environment variable override
if os.environ.get("UNSLOTH_MOE_FORCE_AUTOTUNE", "0") == "1":
force_autotune = True
# 0. Check for Heuristic Override (Skip Autotuning)
# If the GPU is capable (sm_90+), use our "Safe Heuristic"
# This avoids the 2-4 minute startup cost.
if not force_autotune:
# We can try to rely on heuristic
# 0. Check for environment variable override to DISABLE autotuning
if os.environ.get("UNSLOTH_MOE_DISABLE_AUTOTUNE", "0") == "1":
logger.info(
f"Using Heuristic (Safe) MoE kernel configs for SM{device_capability[0]}{device_capability[1]}"
f"UNSLOTH_MOE_DISABLE_AUTOTUNE=1: Using Heuristic (Safe) MoE kernel configs for SM{device_capability[0]}{device_capability[1]}"
)
return _get_heuristic_configs()
if not force_autotune and cache_key in _kernel_config_cache: