feat(forge): support Forge agent and add image_kernel task type#69
Conversation
irvineoy
left a comment
irvineoy
left a comment
There was a problem hiding this comment.
Overall, the direction of this PR looks valuable, especially the Forge integration, image-backed kernel tasks, harness integrity checks, and centralized timing. However, I found several issues that can cause stale or incorrect kernels to be evaluated. I recommend addressing the P1 findings before merging.
| dict(m=128, n=256, k=512, dtype_str="float16"), | ||
| ] | ||
|
|
||
| PERF_CASES = [ |
There was a problem hiding this comment.
[P1] Validate the scored performance shapes for correctness
The new image_kernel runners use separate CASES and PERF_CASES, but run_correctness() only validates CASES. Here, the small correctness shapes differ completely from the 4096/8192 shapes that are actually scored. A kernel can therefore remain correct for the validation shapes while producing incorrect results—or specializing invalid behavior—for the scored shapes and still receive a valid performance score. The same pattern appears in all seven new image-kernel runners.
Please validate every PERF_CASES configuration against its reference at least once before timing it, or include the performance configurations in the correctness suite.
There was a problem hiding this comment.
Fixed. run_correctness() now validates PERF_CASES in addition to CASES. It builds the combined list [*CASES, *scored_cfgs] (where scored_cfgs are the shapes extracted from PERF_CASES) and runs each through the aiter-vs-torch reference check, so every scored shape is verified against its reference before it can earn a performance score. This is applied consistently across all seven image-kernel runners, not just this one.
| return 0 | ||
|
|
||
| task_config_data = _load_task_config(task_config) | ||
| passed, err = evaluate_correctness(Path(workspace), task_config_data) |
There was a problem hiding this comment.
[P1] Recompile each Forge candidate before correctness and benchmarking
The generic adapter calls evaluate_correctness() and measure_performance() directly, without first running evaluate_compilation(). This evaluates stale binaries for tasks where compilation is a separate step. For example, the existing hip2hip/others/matrix_multiplication and mla_decode runners build an executable in compile mode, while their correctness and performance modes only execute the existing binary. Arena compiles the baseline before launching Forge, but source edits made during a Forge iteration do not rebuild that binary.
Consequently, Forge may validate and benchmark the pre-agent binary rather than the edited kernel. Please make each candidate follow a compile → correctness → benchmark sequence, or otherwise guarantee that build artifacts are regenerated after every source change.
There was a problem hiding this comment.
Fixed. Both do_correctness() and do_bench() now call _rebuild_edited_source() first, which invokes Arena's evaluate_compilation() to regenerate build artifacts from the current (agent-edited) source before validating/timing. A failed rebuild aborts the step (correctness reports allclose: False, bench returns a non-zero exit) so a stale pre-edit binary is never validated or benchmarked. JIT-only tasks (no compile_command) skip the step harmlessly. This gives each candidate a compile → correctness → benchmark sequence.
| the API / config); no environment override is needed. | ||
| """ | ||
| model = str(eval_config.get("target_gpu_model", "")).upper() | ||
| return _GPU_ARCH_MAP.get(model, "gfx942") |
There was a problem hiding this comment.
[P1] Do not silently map unsupported GPU models to gfx942
_resolve_gpu_arch() defaults every unrecognized GPU model to gfx942. RDNA4 is supported elsewhere in this repository as gfx1201, but it is absent from this mapping, so a Forge run targeting RDNA4 would be given the MI300 architecture. This can produce invalid build flags, misleading agent guidance, or kernels optimized for the wrong ISA.
Please reuse the repository's shared GPU architecture resolution, use the architecture already detected by setup_rocm_env, or fail explicitly for unknown models instead of silently falling back to gfx942.
There was a problem hiding this comment.
Fixed. _resolve_gpu_arch() no longer contains a hardcoded map defaulting to gfx942. It now reuses Arena's shared resolution — the same priority as src.preprocessing.setup_rocm_env: first the real hardware arch from rocminfo (_detect_gfx_arch_from_rocminfo()), then the configured target_gpu_model looked up in the single source of truth (the architecture map in default_cheatsheet.yaml, which covers RDNA4→gfx1201, MI300/MI325/MI355X, ...). An unknown model now raises ValueError with a clear message instead of silently falling back to gfx942.
| process.wait(timeout=timeout_seconds) | ||
| except subprocess.TimeoutExpired: | ||
| logger.warning(f"Forge loop timed out after {timeout_seconds}s; terminating") | ||
| process.terminate() |
There was a problem hiding this comment.
[P2] Terminate the complete Forge process tree on timeout
The Forge command is launched with shell=True, and this timeout handler terminates only the shell process. The kernel-agents process and its Claude/GPU subprocesses may survive after the shell exits, continuing to consume GPU resources or modify the workspace while Arena runs git checkout and performs final evaluation.
Please launch the command as an argument list with shell=False, create a separate process group/session, and terminate the entire group on timeout.
There was a problem hiding this comment.
Fixed. The command is now launched from the argv list with shell=False and start_new_session=True, so the forge-loop, kernel-agents, and the Claude/GPU subprocesses share one process group led by the child PID. On timeout, _terminate_process_group() signals the whole group (os.killpg with SIGTERM, then SIGKILL after a grace period) rather than just the leader, so no descendants survive to hold the GPU or edit the workspace while Arena runs git checkout and final scoring. (The shlex.quote-joined cmd string is kept only for human-readable logging.)
| joined = " ".join(strs) | ||
|
|
||
| if "aiter" in joined: | ||
| os.environ.setdefault("AITER_REBUILD", "1") |
There was a problem hiding this comment.
[P2] Scope AITER_REBUILD to the current task
This modifies the parent process environment and never restores it. Workers execute multiple tasks in the same process, so an AITER HIP task can leave AITER_REBUILD=1 enabled for subsequent AITER Triton tasks, triggering unnecessary C++ rebuilds and potentially causing large runtime regressions. In addition, setdefault() does not actually force rebuilding if the inherited environment already contains AITER_REBUILD=0.
Please scope this variable to the relevant task subprocesses, or save and restore its previous value when the task finishes. When rebuilding is required, set it explicitly to 1.
There was a problem hiding this comment.
Fixed. AITER_REBUILD is no longer written to os.environ. apply_jit_rebuild() returns the override as a dict ({"AITER_REBUILD": "1"}), and callers in evaluator.py / performance.py apply it only to the specific build subprocess via run_command(..., extra_env=rebuild_env). Since it never touches the parent env, it cannot leak into a later task run by the same worker. It's also set explicitly to "1" (not setdefault) so an inherited AITER_REBUILD=0 is still correctly overridden when a rebuild is required.
jiaqiang-dot-liu
changed the title
…ss hardening Integrate KernelForge's forge-loop as a first-class Arena agent and add the image_kernel task type for optimizing kernels that already ship inside a container image, together with correctness, timing, and anti-gaming hardening. - forge agent: bridge Arena task workspaces to `kernel-agents forge-loop` via a generated driver shim that reuses Arena's own compile/correctness/performance eval; resolve gfx arch through the shared cheatsheet, derive --max-hours from timeout_seconds, and hard-kill the whole forge process group on timeout. - image_kernel task type: seed the in-image source tree directly into each run workspace (reflink copy, no clone/cache, .git dropped); add 7 aiter/sglang tasks (gemm, fp8_gemm, mha_batch_prefill, pa_decode, pa_ragged, rmsnorm, sglang store_cache). - aiter JIT rebuild: centralize forcing kernel edits to take effect (AITER_REBUILD + stale CK .so cleanup), scoped to the build subprocess so it never leaks into later tasks run by the same worker. - harness guard: snapshot and verify task harness/test files around agent runs so kernel scores cannot be gained by editing the measurement harness. - timing: prefer device-side timings and reject host/wall timings, detect and reject empty cuda-graph capture, and share the cuda-graph benchmark helpers across aiter/flydsl/image_kernel tasks. Co-authored-by: Cursor <[email protected]>
e120ba1 to
7ab123a
Compare
2 participants
No description provided.