NVIDIA / NVIDIA/cccl

Reduce `cuda.compute` CI JIT overhead with disk cache or AOT artifacts

Open
#10,530 4 comments 0 reactions 0 assignees View on GitHub
Dominant language
C++
Stars
2.5k
Forks
486
Avg merge
2d 6h
Merged PRs (30d)
295

Description

`cuda.compute` CI recompiles many of the same algorithm specializations across pytest workers and between the non-large and large test passes. We should investigate and implement the highest-impact approach for avoiding repeated JIT work in test jobs.

## Experiment

Ran the same two `cuda.compute` test commands used by CI:
- `pytest -q -n 6 compute/ -m "not large and not free_threading"`
- `pytest -q -n 0 compute/ -m "large and not free_threading"`

Compared current behavior against a local disk-cache prototype:
- Baseline: current in-process cache only.
- Prototype: cache started empty, serialized compiled build results to disk, reused them across pytest workers and between the two pytest invocations.

## Results

| Mode | Non-large | Large | Total |
|---|---:|---:|---:|
| Current behavior | 325.36s | 693.04s | 1018.40s |
| Empty disk cache at job start | 305.80s | 211.63s | 517.43s |

Result: **~17.0 min -> ~8.6 min**, about **49% faster**.

Cache stats:
- Non-large: 81 hits, 739 misses
- Large: 89 hits, 20 misses
- Final cache size: 751 blobs, 196 MB

## Options To Investigate

1. Add disk caching in the test job.
- Cache serialized `cuda.compute` build results on disk.
- Share cache across pytest workers and between non-large/large invocations.
- Lower implementation cost.
- Still pays first-use JIT cost inside each test job.

2. Precompile test artifacts during the wheel build.
- Explicitly AOT-compile known test algorithm specializations during `build_cuda_cccl_python`.
- Include serialized artifacts in the wheel or as a wheel-adjacent artifact.
- Test jobs load precompiled artifacts and avoid most/no JIT.
- Potentially larger CI win, higher implementation cost.

A staged approach may make sense:
- Stage 1: add safe disk caching and CI hit/miss stats.
- Stage 2: use the stats to identify high-value specializations for AOT precompile.
- Stage 3: move those artifacts into the wheel build path if the extra build cost is worth the test-time savings.

## Done When
- We have measured disk-cache vs AOT-precompile impact.
- The selected approach is enabled for `test_cuda_compute_python` CI jobs.

Contributor guide

Open the contributing guide

Research direction

Start by running the two listed pytest commands for compute/ and compare the current cache behavior with the disk-cache prototype and its hit/miss results. Read the build_cuda_cccl_python path before assessing AOT artifacts. Done means measuring both approaches and enabling the selected one for test_cuda_compute_python CI jobs.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
ci-cd, performance, testing-qa
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.