auto-comm-kernels

Workspace repo for studying ParallelKernelBench (PKB), cloned locally at ParallelKernelBench/ (git-ignored here since it's tracked by its own upstream history — pull it directly with git -C ParallelKernelBench pull), plus the papers behind the effort to apply Hawkeye's taxonomy-of-unit-tests method to hillclimb it — see papers/.

Papers

• papers/parallelkernelbench.pdf — Chan, Paek, Guo, Arora, Fu. ParallelKernelBench: Can LLMs Write Fast Multi-GPU Kernels? Stanford/Together AI/UCSD, 2026.
• papers/hawkeye.pdf — Tschand, Ramakrishnan, Ingare, Guo, Ma, Wan, Arora, Mirhoseini, Reddi. Hawkeye: Hardware-Aware GPU Kernel Optimization with Minimal Supervision. Harvard/Stanford/Together AI/Caltech, 2026.
• papers/applying-hawkeye-to-pkb.md — analysis of both papers against the actual PKB codebase: how Hawkeye's taxonomy method maps onto PKB's existing agentic harness (kernelgen/generate_kernel_agent.py), the proposed taxonomy (7 rows, no columns — PKB's own Appendix G rules out architecture as a column axis), what's missing in PKB's current prompt structure per Hawkeye's own ablations, and where the analogy strains.
• papers/taxonomy-study-guide.md — self-contained per-row study packets (context + real code + questions) for learning the concepts behind each of the 7 taxonomy rows, meant to be used in a separate chat.

What PKB is

A benchmark from Together AI (paper: ParallelKernelBench: Can LLMs Write Fast Multi-GPU Kernels?, Chan/Paek/Guo/Arora/Fu, 2026) testing whether LLMs can turn PyTorch + NCCL reference implementations of multi-GPU operations into fast, correct low-level kernels (CUDA, Triton, or ParallelKittens/ThunderKittens). Design is explicitly modeled on Stanford's KernelBench, extended to the multi-GPU/distributed setting.

Repo layout

• reference/ — 87 numbered problems (1_allreduce.py ... 87_conv2d_boundary_exchange.py), each a solution(...) function written in plain PyTorch + torch.distributed/NCCL. Problems span:
  • Core collectives: allreduce, allgather, broadcast, reduce, scatter, reducescatter, alltoall, gather (problems 1–8)
  • Fused comm+compute: allgather-gemm, gemm-reducescatter, RoPE+allgather, TP RMSNorm, fp8 quant/dequant blocks
  • MoE: token preprocessing, all-to-all primitives, expert-parallel GEMMs (balanced/wide/narrow), fused MoE forward (+LoRA)
  • Sequence/context parallelism: Ulysses-style all-to-all attention sharding, ring attention (fwd/bwd, combined with TP/PP)
  • Training infra: DDP, ZeRO-1/2 shard, FSDP step end-to-end, FSDP+TP, fused Adam with grad unshard/allgather, gradient clipping (with/without expert parallel)
  • Domain-specific: OpenCLIP contrastive loss, distributed FFT (PhysicsNeMo), 3D Gaussian splatting (gsplat), spherical convolution, GNN neighbor sampling / feature exchange, TorchRec KJT all2all, Hyena conv boundary exchange, vocab-parallel cross-entropy/log-prob top-k, DINOv2 distributed KNN/Sinkhorn-Knopp, SAM3 allgather+NMS
• kernelgen/ — prompt assembly (prompts.toml) + generation drivers: generate_kernel.py (single-shot LLM call against ALLOWED_MODELS: Together/Gemini/Anthropic/OpenAI) and generate_kernel_agent.py (agentic loop via mini-swe-agent, iterating dryrun→eval against run_local.py).
• scripts/worker.py — the actual per-rank execution logic launched under torchrun. Loads a solution fn from a candidate file, runs backend-specific init (reference/triton/cuda/parallelkittens), executes, saves per-rank tensors, optionally times (500 warmup / 100 measured iters) and profiles.
• run_local.py / run_modal.py / run_together.py — three ways to drive the harness: bare-metal multi-GPU node, Modal, or Together's cloud container service (together beta jig). All wrap dryrun (run one backend, no comparison) and eval (run reference + candidate on identical inputs across --trials seeds, diff per-rank .pt outputs within --atol/--rtol).
• utils/ — shared plumbing: input_output_tensors.py (1888 lines — per-problem shape/dtype/RNG input construction, the largest file in the repo), compare_outputs.py / compare_performance.py (offline artifact inspection), init_and_finalize_backends.py, cuda_helpers.py, triton_utils.py, parallelkittens_runtime.py, performance.py.

How evaluation works

1. Generate: LLM writes solutions_<backend>_<precision>_<hardware>_<provider>_<model>/{stem}_{backend}.py.
2. Dryrun: torchrun --nproc-per-node N run_local.py --mode dryrun runs just the candidate to confirm it compiles/runs/produces rank_*.pt.
3. Eval: same harness runs reference then candidate on matching random inputs across multiple trials, diffs per-rank tensors, stops on first mismatch outside tolerance.
4. Perf (optional): 500 warmup + 100 timed iterations, speedup vs. reference reported per rank.

Stack

• Python ≥3.12, uv-managed (pyproject.toml + uv.lock), PyTorch ≥2.9 (cu128 index), Triton, cuda-core, mpi4py, ThunderKittens (cloned separately, THUNDERKITTENS_ROOT) for the ParallelKittens backend.
• Targets Hopper (H100) by default (TORCH_CUDA_ARCH_LIST=9.0); cloud eval path builds a container image ([tool.jig.image]) deployed at 8×H100.
• License: Apache 2.0. Actively developed — 39 commits, most recent work adding new problem categories (physics/GNN/vision-language problems) and defaulting eval to bf16.

Notable design choices

• Correctness is checked per-rank, not just on a gathered/rank-0 result — candidates must match the reference's exact output-sharding behavior, not just produce a numerically-equivalent final answer.
• Perf timing follows the ThunderKittens 2 blog's benchmarking methodology (fixed warmup/measure iteration counts via CUDA events).
• Adding a new backend requires wiring init_<backend>/finalize_<backend> into scripts/worker.py and updating the prompt template — the harness assumes one .py file per solution regardless of backend.