SIGN IN SIGN UP

[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