[Enhancement] Speed up cold parallel/AOT compilation up to ~4x (#2809)
* [Enhancement] Speed up cold parallel/AOT compilation up to ~4x Three independent, low-risk changes (measured 45.6s -> 11.5s for 144 cold kernels on an H20, 180 cores): 1. cache: move kernel disk-save out of the global lock. KernelCache.cached() held the class lock around _save_kernel_to_disk (source + .so copy + cloudpickle), serializing every worker's save. The save stages+renames atomically and is idempotent, so it is already concurrency-safe without the lock (the disk load was already outside it). Lock now guards only _memory_cache. This is the bulk of the win (~3.1x). 2. jit: core-scale par_compile workers. Default was min(32, cpu+4); now min(len(funcs), available_cpus) when unset, with TILELANG_PAR_COMPILE_WORKERS override. Lowering is GIL-releasing C++ and nvcc is a subprocess, so threads parallelize. get_available_cpu_count is now cgroup-aware: it respects cpuset affinity AND caps at the cgroup CFS bandwidth quota (v2 cpu.max / v1 cpu.cfs_quota_us,period_us), so a CPU-throttled container does not oversubscribe nvcc jobs. 3. nvcc: opt-in parallel device compile (off by default). TL_NVCC_THREADS=N -> --threads (CUDA >= 11.2) + --split-compile (>= 12.1); scheduling-only, identical SASS. Helps large multi-kernel TUs; no effect on single-kernel TUs. get_available_cpu_count moved to tilelang.utils.device to break the autotuner<->jit import cycle; unit-tested in testing/python/utils. Adds benchmark/compile_speed/ to reproduce the numbers: a zoo of realistic inference kernels (GEMM, GQA attention, RMSNorm, SwiGLU, softmax) over Qwen2.5/Llama-3 dims (~126 distinct cold kernels at the default scale), with a before/after that reconstructs the pre-change baseline (lock-save + 32 workers) in one command (3.1x on the H20). * Simplify parallel compilation controls --------- Co-authored-by: SiriusNEO <chaofan@deepseek.com>
C
cklxx committed
478ab70d8dd1dfea3416f73a43dcbc6bb9cc2be3
Parent: a426ff3
Committed by GitHub <noreply@github.com>
on 8/3/2026, 7:55:31 AM