diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index fda66255c..6aea6d8b9 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -113,6 +113,7 @@ SGLang supports various environment variables that can be used to configure its | --- | --- | --- | | `SGLANG_WAIT_WEIGHTS_READY_TIMEOUT` | Timeout period for waiting on weights | `120` | | `SGLANG_DISABLE_OUTLINES_DISK_CACHE` | Disable Outlines disk cache | `true` | +| `SGLANG_USE_CUSTOM_TRITON_KERNEL_CACHE` | Use SGLang's custom Triton kernel cache implementation for lower overheads (automatically enabled on CUDA) | `false` | ## Function Calling / Tool Use diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 0ac752b93..05274e558 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -3618,6 +3618,12 @@ def cached_triton_kernel(key_fn=None): """ def decorator(fn): + # Auto-enable the custom kernel cache for CUDA, where it is + # known to be compatible. + if is_cuda() and not envs.SGLANG_USE_CUSTOM_TRITON_KERNEL_CACHE.is_set(): + logger.debug("Detected platform CUDA, using custom triton kernel cache.") + return CachedKernel(fn, key_fn) + if envs.SGLANG_USE_CUSTOM_TRITON_KERNEL_CACHE.get(): logger.debug( f"{envs.SGLANG_USE_CUSTOM_TRITON_KERNEL_CACHE.name} = True. Using custom triton kernel cache."