Scale beam search: multi-LLM, samples-per-prompt, multi-GPU, PTX dedup, NCU caching

[jiannanWang]

Summary

Beam search on main runs with a single LLM, a single GPU, and no candidate deduplication. That's enough for easy problems but plateaus quickly: a single model produces near-duplicate kernels each round, and on harder problems a single model often can't produce any working kernel at all (e.g., claude-opus-4.5 alone yields 0/32 successful kernels on the Gemma3 fused gate-up workload).

This PR scales beam search to run multiple LLMs in parallel, expand multiple samples per (parent, bottleneck, model) triple, distribute workers across multiple GPUs, and deduplicate beam entries by compiled-PTX fingerprint so the beam doesn't fill with byte-equivalent copies of the leader.

All new features are opt-in via strategy_config. Existing flows that don't set models or samples_per_prompt are unchanged.

Example commands

Three new presets ship with the PR:

# 90-worker concentrated: 2 parents × 3 bottlenecks × 3 models × 5 samples                                                                                                   
python examples/run_opt_manager.py \                                                                                                                                         
  --kernel-dir examples/optimize_01_matvec \                                                                                                                                 
  --strategy beam_search_diverse_concentrated \                                                                                                                              
  --max-rounds 8                                                                                                                                                             
                                                                                                                                                                             
# 90-worker spread: 5 parents × 3 bottlenecks × 3 models × 2 samples                                                                                                         
python examples/run_opt_manager.py \     
  --kernel-dir examples/optimize_01_matvec \
  --strategy beam_search_diverse \       
  --max-rounds 8       

# 9-worker smoke for pipeline validation 
python examples/run_opt_manager.py \                             
  --kernel-dir examples/optimize_01_matvec \
  --strategy beam_search_diverse_smoke \ 
  --max-rounds 1

Each preset configures models: [claude-opus-4.6, gpt-5-4, gemini-2-5-pro] out of the box, routed via the existing RelayProvider fallback.

Experiment results

Two problems, 8 rounds each:

┌───────────────┬───────────────────────────────────┬──────────────────────────────────────────────┐                                                                         
│    Problem    │ main (4 workers, claude-opus-4.5) │  This PR (90 workers, 3 LLMs concentrated)   │
├───────────────┼───────────────────────────────────┼──────────────────────────────────────────────┤                                                                         
│ matvec        │ 2.1494 ms — 0.95× PyTorch eager   │ 1.8318 ms — 1.12× PyTorch eager              │
├───────────────┼───────────────────────────────────┼──────────────────────────────────────────────┤
│ gemma3_swiglu │ failed: 0/32 working kernels      │ 0.2716 ms — 3.10× eager, 1.14× torch.compile │                                                                         
└───────────────┴───────────────────────────────────┴──────────────────────────────────────────────┘                                                                         

The matvec improvement comes from richer fanout exposing more refinement directions per round. The gemma3_swiglu improvement comes from multi-LLM breaking through the single-model correctness floor — gpt-5-4 and gemini-2-5-pro find working kernels for this problem where claude-opus-4.5 does not.

What is changed

• models: list[str] strategy knob — per-candidate LLM override; each worker round-robins through the configured list.
• samples_per_prompt: int strategy knob — N independent LLM draws per (parent, bottleneck, model) triple. Default = 1.
• Per-GPU benchmark lock pool — dict[gpu_id, mp.Lock] replaces the single global lock; each worker acquires only its own GPU's lock.
• Worker→GPU pinning — CUDA_VISIBLE_DEVICES=<gpu_id> is set in each worker process before any torch import.
• nvidia-smi-based GPU detection — OptimizationManager._detect_gpus no longer touches torch.cuda (which would lock the device list for forked children).
• PTX-fingerprint module — searching/ptx_fingerprint.py normalizes PTX (strip comments, canonicalize registers / labels) and SHA-256-hashes it; persisted on ProgramEntry.ptx_hash.
• PTX capture in benchmark — sets TRITON_CACHE_DIR= and reads the resulting *.ptx files to compute the hash alongside time_ms.
• PTX-hash dedup in beam selection — BeamSearchStrategy.update_with_results collapses pool entries with identical hash to the fastest representative before sort+truncate. New PTX dedup X→Y line in round logs.
• Per-parent NCU caching — manager profiles each unique parent kernel once per round and shares the result with siblings instead of N re-profiles.
• Three new strategy presets — beam_search_diverse.yaml, beam_search_diverse_concentrated.yaml, beam_search_diverse_smoke.yaml. Existing presets unchanged.