Add parallel CuTeDSL tuning - #246
Conversation
Human Note Agent note A CUDA-owning process cannot safely fork workers that rediscover the device or inherit CuTeDSL compiler state. Add a fresh interpreter driver that receives only static, pickleable compile calls and parent-discovered target metadata, imports the cached compile function once, then forks bounded workers to populate distinct disk artifacts. The parent loads those artifacts and benchmarks each candidate sequentially with Inductor's CUDA benchmarker. Expose `run_tunable` as the author-facing boundary. An op owns a conservative default, an input-aware typed config generator, a projection from runtime tensors to positional static compile arguments, and a launch adapter. Runtime tensors never cross the compiler-process boundary, and a fully warm tune starts no compiler process. Test Plan: ```bash ~/.venvs/nightly/bin/python -m pytest -q test/test_cute_cache.py gpu-run 0 -- ~/.venvs/nightly/bin/python -m pytest -q test/test_cute_cache_gpu.py ~/.venvs/dev/bin/ruff check attn_gym/_backends/cute test/test_cute_cache.py test/test_cute_cache_gpu.py ``` stack-info: PR: #246, branch: drisspg/stack/39
24fd863 to
0dd6382
Compare
ede065d to
986dcb6
Compare
meta-cla
Bot
added
the
CLA Signed
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0dd6382c76
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| target=target, | ||
| timeout=timeout, | ||
| ) | ||
| return [fn(*args) for args in call_args] |
There was a problem hiding this comment.
When two host threads call compile_many() or tune() concurrently with different explicit targets, each call sets the process-global target before blocking in its compiler subprocess, so the later call can overwrite it before the first reaches this load loop. The first call then keys fn(*args) using the other thread's target, potentially returning a kernel for the wrong GPU architecture or compiling that variant in the CUDA-owning parent; target selection must remain scoped or synchronized through the complete artifact-loading phase.
Useful? React with 👍 / 👎.
| text=True, | ||
| env={**os.environ, "TORCH_WARM_POOL": "0"}, | ||
| timeout=timeout, |
There was a problem hiding this comment.
If compilation exceeds timeout, subprocess.run kills only the fresh driver process; the forked ProcessPoolExecutor workers are grandchildren and remain orphaned, continuing expensive compilation and cache writes after precompile_many() has raised. Launch the driver in a managed process group and terminate the entire group on timeout so the advertised bound also applies to its workers.
Useful? React with 👍 / 👎.
0dd6382 to
9a34d18
Compare
1 participant
Stacked PRs:
Add parallel CuTeDSL tuning
Human Note
Agent note
A CUDA-owning process cannot safely fork workers that rediscover the device or inherit CuTeDSL
compiler state. Add a fresh interpreter driver that receives only static, pickleable compile calls
and parent-discovered target metadata, imports the cached compile function once, then forks bounded
workers to populate distinct disk artifacts. The parent loads those artifacts and benchmarks each
candidate sequentially with Inductor's CUDA benchmarker.
Expose
run_tunableas the author-facing boundary. An op owns a conservative default, aninput-aware typed config generator, a projection from runtime tensors to positional static compile
arguments, and a launch adapter. Runtime tensors never cross the compiler-process boundary, and a
fully warm tune starts no compiler process.
Test Plan: