[Enhancement] Speed up cold parallel/AOT compilation up to ~4x - #2809
Conversation
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughChangesCPU-aware worker selection is centralized, CUDA compilation gains optional parallelism and ccache support, kernel cache persistence moves outside the lock, and a benchmark with generated kernel workloads measures cold parallel-compilation throughput. Compilation and cache updates
Estimated code review effort: 3 (Moderate) | ~30 minutes Sequence Diagram(s)sequenceDiagram
participant par_compile
participant KernelCache
participant compile_cuda
participant ccache
participant nvcc
par_compile->>KernelCache: compile cold kernel
KernelCache->>compile_cuda: invoke CUDA compilation
compile_cuda->>ccache: optionally wrap command
ccache->>nvcc: forward NVCC flags and source
nvcc-->>compile_cuda: compiled artifact
compile_cuda-->>KernelCache: compilation result
KernelCache->>KernelCache: save artifact without holding lock
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tilelang/cache/kernel_cache.py`:
- Around line 404-415: Coordinate cache persistence with clear_cache so clearing
is mutually exclusive with in-flight saves. Update the cache flow around
_save_kernel_to_disk and memory-entry publication to register active saves
before disk persistence, block new persistence while clear_cache() runs, and
have clear_cache() wait for all active saves to finish before removing state.
Ensure a save cannot republish its memory entry after clearing completes.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: f839f30e-b381-4f37-80d0-fc319ffbb78e
📒 Files selected for processing (6)
tilelang/autotuner/tuner.pytilelang/cache/kernel_cache.pytilelang/contrib/nvcc.pytilelang/jit/__init__.pytilelang/jit/adapter/libgen.pytilelang/utils/device.py
f96178d to
40a0fed
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tilelang/utils/device.py`:
- Around line 26-38: Update get_available_cpu_count to detect cgroup CPU
bandwidth quotas from v2 cpu.max and v1 cpu.cfs_quota_us/cpu.cfs_period_us, and
cap the affinity or os.cpu_count result at ceil(quota / period) when a positive
quota is configured. Preserve the existing affinity-first and fallback behavior
for quota-less or unavailable cgroup environments, and continue returning at
least 1.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: b1457868-a114-444f-8f20-215fa3fdd796
📒 Files selected for processing (6)
tilelang/autotuner/tuner.pytilelang/cache/kernel_cache.pytilelang/contrib/nvcc.pytilelang/jit/__init__.pytilelang/jit/adapter/libgen.pytilelang/utils/device.py
40a0fed to
af78d2d
Compare
cklxx
changed the title
af78d2d to
9c80d6b
Compare
9c80d6b to
0b2abce
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@benchmark/compile_speed/README.md`:
- Around line 44-46: Remove the duplicated consecutive “126 cold kernels, H20 /
180 cores:” heading in the benchmark README, leaving a single occurrence.
In `@tilelang/contrib/nvcc.py`:
- Around line 101-122: Update the TL_NVCC_THREADS validation in the
environment-variable parsing flow to treat zero as disabled, returning no flags
for n == 0. Keep positive values enabled and preserve the existing
negative-value rejection and CUDA-version gating in the surrounding flag
construction.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 357845c0-957b-4607-9e79-bf96c37f1790
📒 Files selected for processing (10)
benchmark/compile_speed/README.mdbenchmark/compile_speed/benchmark_compile_speed.pybenchmark/compile_speed/kernel_zoo.pytesting/python/utils/test_device_cpu_count.pytilelang/autotuner/tuner.pytilelang/cache/kernel_cache.pytilelang/contrib/nvcc.pytilelang/jit/__init__.pytilelang/jit/adapter/libgen.pytilelang/utils/device.py
| raw = os.environ.get("TL_NVCC_THREADS") | ||
| if raw is None: | ||
| return [] | ||
| try: | ||
| n = int(raw) | ||
| except ValueError: | ||
| return [] | ||
| if n < 0: | ||
| return [] | ||
|
|
||
| try: | ||
| version = get_cuda_version() # (major, minor, ...) | ||
| major, minor = version[0], version[1] if len(version) > 1 else 0 | ||
| except Exception: | ||
| return [] | ||
|
|
||
| flags: list[str] = [] | ||
| if (major, minor) >= (11, 2): | ||
| flags += ["--threads", str(n)] | ||
| # --split-compile was introduced in CUDA 12.1 (absent from 11.x and 12.0 nvcc). | ||
| if (major, minor) >= (12, 1): | ||
| flags += ["--split-compile", str(n)] |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C3 'TL_NVCC_THREADS|if n < 0|--threads|--split-compile' tilelang/contrib/nvcc.py
nvcc --help | grep -A6 -- '--threads'Repository: tile-ai/tilelang
Length of output: 1069
🌐 Web query:
NVIDIA nvcc --threads --split-compile zero value use all CPUs documentation
💡 Result:
According to official NVIDIA CUDA documentation, the --split-compile option (and its extended variant, --split-compile-extended) allows you to perform compiler optimizations in parallel to reduce compile time [1][2]. When you set these options to 0 (e.g., --split-compile=0 or --split-compile-extended=0), the compiler is instructed to use the maximum number of threads available on your system [1][3]. Key details regarding these flags include: - Functionality: These options split device code into smaller translation units, allowing optimization passes to run concurrently across multiple threads before linking them back together prior to code generation [1][4]. - --split-compile=1: Setting the value to 1 causes the option to be ignored [1][2]. - Compatibility: --split-compile can be used in conjunction with device Link Time Optimization (-dlto) and the --threads flag [1][3]. --split-compile-extended is only applicable when using -dlto [1][2]. - Trade-offs: While split compilation is intended to accelerate compilation, --split-compile-extended is noted as a more aggressive form that may potentially impact the performance of the resulting compiled binary [1][3].
Citations:
- 1: https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/
- 2: https://docs.nvidia.com/cuda/archive/12.8.0/cuda-compiler-driver-nvcc/index.html
- 3: https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/index.html
- 4: https://docs.nvidia.com/cuda/archive/12.6.3/cuda-compiler-driver-nvcc/index.html
🌐 Web query:
site:docs.nvidia.com/cuda nvcc --threads N 0 use number of threads
💡 Result:
The --threads option (abbreviated as -t) in the NVIDIA CUDA Compiler Driver (nvcc) is used to specify the maximum number of CPU threads used to execute compilation steps in parallel [1]. When you use the argument 0 with this option (i.e., --threads 0 or -t 0), the compiler defaults to using a number of threads equal to the total number of CPUs available on the machine [1]. Key details regarding this option: - Purpose: It is designed to improve compilation speed, particularly when compiling a single compilation unit for multiple GPU architectures [2][1]. - Behavior: If the specified number is 1, the option is ignored [1]. If the number is greater than 1, the compiler creates that specific number of threads to perform compilation steps in parallel [1]. - Compatibility: This option can work in conjunction with other compilation features such as --split-compile and --split-compile-extended, which also use thread-based parallelism to optimize the compilation process [1].
Citations:
- 1: https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/index.html
- 2: https://docs.nvidia.com/cuda/cuda-programming-guide/02-basics/nvcc.html
🌐 Web query:
site:docs.nvidia.com/cuda nvcc --threads NVCC_THREADS threads option 1 ignored default thread count
💡 Result:
The --threads option (abbreviated as -t) in nvcc is used to specify the number of CPU threads utilized to parallelize the compilation of a single compilation unit for multiple GPU architectures [1][2]. Regarding your specific query about the option being ignored when set to 1, official NVIDIA documentation historically confirms this behavior: if the number specified is 1, the --threads option is ignored [3]. This is because a thread count of 1 provides no parallelism, so the compiler defaults to single-threaded execution [3]. If the number is set to 0, nvcc will typically use the number of available CPU threads [3].
Citations:
- 1: https://docs.nvidia.com/cuda/cuda-programming-guide/02-basics/nvcc.html.md
- 2: https://docs.nvidia.com/cuda/archive/13.2.0/cuda-programming-guide/02-basics/nvcc.html
- 3: https://docs.nvidia.com/cuda/archive/11.2.2/pdf/CUDA_Compiler_Driver_NVCC.pdf
Treat TL_NVCC_THREADS=0 as disabled.
NVCC interprets 0 as “use all available CPU threads” for both --threads and --split-compile, so parallel AOT can overflow the cgroup-aware worker limits. Reject or ignore zero instead of emitting the flags.
Proposed fix
- if n < 0:
+ if n <= 0:
return []📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| raw = os.environ.get("TL_NVCC_THREADS") | |
| if raw is None: | |
| return [] | |
| try: | |
| n = int(raw) | |
| except ValueError: | |
| return [] | |
| if n < 0: | |
| return [] | |
| try: | |
| version = get_cuda_version() # (major, minor, ...) | |
| major, minor = version[0], version[1] if len(version) > 1 else 0 | |
| except Exception: | |
| return [] | |
| flags: list[str] = [] | |
| if (major, minor) >= (11, 2): | |
| flags += ["--threads", str(n)] | |
| # --split-compile was introduced in CUDA 12.1 (absent from 11.x and 12.0 nvcc). | |
| if (major, minor) >= (12, 1): | |
| flags += ["--split-compile", str(n)] | |
| raw = os.environ.get("TL_NVCC_THREADS") | |
| if raw is None: | |
| return [] | |
| try: | |
| n = int(raw) | |
| except ValueError: | |
| return [] | |
| if n <= 0: | |
| return [] | |
| try: | |
| version = get_cuda_version() # (major, minor, ...) | |
| major, minor = version[0], version[1] if len(version) > 1 else 0 | |
| except Exception: | |
| return [] | |
| flags: list[str] = [] | |
| if (major, minor) >= (11, 2): | |
| flags += ["--threads", str(n)] | |
| # --split-compile was introduced in CUDA 12.1 (absent from 11.x and 12.0 nvcc). | |
| if (major, minor) >= (12, 1): | |
| flags += ["--split-compile", str(n)] |
🧰 Tools
🪛 Ruff (0.16.0)
[warning] 114-114: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tilelang/contrib/nvcc.py` around lines 101 - 122, Update the TL_NVCC_THREADS
validation in the environment-variable parsing flow to treat zero as disabled,
returning no flags for n == 0. Keep positive values enabled and preserve the
existing negative-value rejection and CUDA-version gating in the surrounding
flag construction.
0b2abce to
1638ee4
Compare
305134d to
f307e31
Compare
|
Pushed
Ready for another look when you have a moment. |
f307e31 to
6ada648
Compare
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).
6ada648 to
8627e8d
Compare
SiriusNEO
left a comment
SiriusNEO
left a comment
There was a problem hiding this comment.
LGTM now, but I think the main gain is from removing the lock (which is useless now).
2 participants
Speeds up cold parallel/AOT compilation. Three independent, low-risk changes:
KernelCachelock. The save stages+renames atomically and is idempotent, so it was already safe without the lock; holding it serialized every worker's save. (the bulk of the win)par_compiledefaults tomin(len(funcs), available_cpus)instead ofmin(32, cpu+4).get_available_cpu_countis cgroup-aware (cpuset affinity + CFS quota).TL_NVCC_THREADS(--threads≥11.2,--split-compile≥12.1), off by default. Scheduling-only (identical SASS); helps large multi-kernel TUs.Measured on H20 / 180 cores, 126 cold kernels (
benchmark/compile_speed, a zoo of GEMM/attention/RMSNorm/SwiGLU/softmax over Qwen2.5/Llama-3 dims):No API/ABI change; #1 and #2 are safe defaults, #3 is opt-in.