Skip to content

kevinchiu19/evokernel

Repository files navigation

EvoKernel — Reproduction Skeleton

A clean, dependency-free Python reproduction of EvoKernel (ICLR 2026): Towards Cold-Start Drafting and Continual Refining: A Value-Driven Memory Approach with Application to NPU Kernel Synthesis.

See EvoKernel_分析.md for the methodology breakdown.


1. What's in here

.
├── EvoKernel_分析.md          # Detailed Chinese analysis of the paper
├── evokernel/                 # Core package
│   ├── __init__.py
│   ├── config.py              # EvoKernelConfig dataclass
│   ├── mdp.py                 # State, Action, Outcome, Trajectory
│   ├── memory.py              # MemoryBank, MemoryItem, QTable
│   ├── retrieval.py           # ValueDrivenRetriever (top-K + Q filter)
│   ├── popart.py              # PopArt-style online normalization
│   ├── verifier.py            # MultiGateVerifier interface + stubs
│   ├── anti_hacking.py        # Rule + LLM auditor
│   ├── generator.py           # LLMGenerator wrapper + prompt assembly
│   └── agent.py               # EvoKernelAgent — solve() loop
├── run_evokernel.py           # End-to-end runner (CLI)
├── tasks_example.json         # Two toy operator specs
├── seeds_example.jsonl        # Initial M_0 knowledge
├── requirements.txt
└── README.md                  # this file

The package is pure-stdlib so you can sanity-check the loop without any LLM/backend installed.


2. The 1-minute mental model

EvoKernel = LLM generator G_θ + a learnable retrieval policy μ over a self-evolving memory bank 𝓜.

For each operator:

state = (task, has_feasible=False)
for t in 1..30:
    if not has_feasible:                       # Stage 1 — Cold-Start Drafting
        c_t  = retrieve_draft(s_t)             # ε-greedy on Q_1
        a_t  = G(prompt(task, c_t))
        o_t  = Verifier(a_t)                   # 4 gates
        r    = +1 if feasible else -1          # Eq. 5
        Q_1[s, m] += α (r − Q_1[s, m])  ∀m∈c_t # MC update, Eq. 4
    else:                                       # Stage 2 — Continual Refining
        p_t  = pick_start_point(P(x))          # ε-greedy on Q_2
        c_t  = retrieve_refine(s_t, profile)
        a_t  = G(prompt(task, p_t, c_t))
        o_t  = Verifier(a_t)
        r    = -1 if !feas else tanh(log b - log ℓ)   # Eq. 6
        r̂    = PopArt(r)                       # online z-score
        Q_2[s, m] += α (r̂ − Q_2[s, m])  ∀m∈c_t∪{p_t}
    memory.append(a_t, o_t)

That entire loop lives in evokernel/agent.py.


3. Quick start (offline / stub mode)

python run_evokernel.py \
    --tasks tasks_example.json \
    --seeds seeds_example.jsonl \
    --model stub \
    --max-iters 6 \
    --out runs/results.json

This runs the full agent loop with a deterministic stub LLM and stub verifier (everything compiles + is "correct" + 1ms latency). You should see trajectories printed for both example operators and a results.json summary.


4. Wiring up real backends

Three integration points — all single-function plugins, no monkeypatching:

Component Where What you provide
LLM make_generator() in run_evokernel.py a (prompt: str) -> str function
Verifier make_verifier() in run_evokernel.py implementations of Compiler, Runner, Profiler (see evokernel/verifier.py)
Dense retriever ValueDrivenRetriever(dense_scorer=...) optional embedding-based scorer to replace lexical bag-of-words

4.1 LLM example (OpenAI)

export OPENAI_API_KEY=sk-...
python run_evokernel.py --tasks tasks_example.json --model openai:gpt-4o

4.2 LLM example (Anthropic)

export ANTHROPIC_API_KEY=sk-ant-...
python run_evokernel.py --tasks tasks_example.json --model anthropic:claude-opus-4-6

4.3 Real Ascend C verifier

Implement the three Protocol classes in evokernel/verifier.py:

class AscendCCompiler:
    def compile(self, sources: dict[str, str]) -> tuple[bool, str]:
        # write sources to a tmp project, invoke Ascend C toolchain, capture log
        ...

class AscendCRunner:
    def run_and_check(self, sources, task_spec, atol, rtol):
        # build the .so, dispatch via torch_npu, compare against the PyTorch ref
        ...

class AscendCProfiler:
    def latency_ms(self, sources, task_spec):
        # invoke msprof three times (PipeUtilization/Memory/ResourceConflict)
        # return mean of step_trace_time.csv "Computing" column
        ...

Then swap them into make_verifier().

CUDA reproduction is symmetric — replace the toolchain with nvcc + Nsight Compute and the rest of the framework is unchanged.


5. Mapping back to the paper

Paper § Equation Code
§3.1 — M-MDP definition mdp.py
§3.2 — heterogeneous M memory.MemoryBank (5 kinds)
§3.2 — Q_1 / Q_2 stage values Eq. 4 memory.QTable
§3.2 — top-K dense recall, λ over-retrieve retrieval.ValueDrivenRetriever._retrieve
§3.3 — drafting reward Eq. 5 agent._draft_step
§3.4 — refining reward Eq. 6 agent._refine_step
§3.4 — PopArt online normalization popart.PopArtNormalizer
§3.5 — multi-gate verifier verifier.MultiGateVerifier
Appendix C — anti-hacking anti_hacking.RuleAntiHack + LLMAntiHack
Appendix B — convergence guarantees Lemmas 1–4 enforced by reward_clip, popart_eps, alpha

6. Reproducing the paper's experiments (sketch)

  1. Build the operator set — clone KernelBench L1 + L2, write each as a task JSON entry (task_id, name, category, level, ref_src, metadata, optional torch_latency_ms).
  2. Seed M_0 — collect Ascend C API templates and best practices into a JSONL like seeds_example.jsonl.
  3. Run the three baselines for comparison:
    • Pass@k: set --max-iters 30 and disable retrieval (run a dummy RetrieverNoOp that always returns []).
    • Refinement: restrict memory to a single task by clearing mem.items between operators.
    • EvoKernel (ours): as configured in run_evokernel.py.
  4. Transfer experiments: run L1 first, persist memory with agent.save(), then start a fresh agent on L2 with --load.
  5. Cross-backbone transfer: build memory.jsonl with model A, then run model B with --load.
  6. Use summarize() in run_evokernel.py to compute pass@T and median speedup. Combine with backend-native profiling for absolute latency.

7. Things you'll want to tune

Knob Default Effect
max_iterations 30 T budget per operator
over_retrieval_lambda 15 larger pool → discovers high-value items but adds noise (paper: 2 → 15 lifted Acc 67% → 84%)
epsilon / epsilon_decay 0.2 / 0.995 exploration in retrieval and start-point selection
alpha 0.3 MC step size for Q updates
popart_warmup 5 wait until N samples before normalizing
enable_llm_anti_hack False enable the GPT-based auditor (cost vs precision trade-off)

8. License & citation

This skeleton is a third-party reimplementation. If you use it, please cite the original paper.

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages