jiliu1/avo_test_cases

k04_fmha_prefill — GEAK task layout

Copy of k04_fmha_prefill/ as avo_data_examples/k04_fmha_prefill_geak/, same top-level GEAK layout as paged_attention_large.

Changes from the original tree

1. Vendored _bench_common + patched kernel/test_harness.py imports (sys.path.append → task scripts/).
2. kernel/.git omitted from the copy (cleaner worktree for AVO).
3. kernel/.rocprofv3/ omitted.

Dependencies

• PyTorch (ROCm)
• aiter (golden flash_attn_func in harness)
• Triton

Commands

cd avo_data_examples/k04_fmha_prefill_geak
python3 scripts/task_runner.py compile
python3 scripts/task_runner.py correctness
python3 scripts/task_runner.py performance

geak-avo

export GEAK_SUBAGENTS_ROOT=/mnt/raid0/models/avo/avo_workspace/GEAK/subagents/preprocess
geak-avo --repo avo_data_examples/k04_fmha_prefill_geak \
  --task "Optimize MLA FMHA prefill _attn_fwd. Metric: latency (lower is better). kernel/kernel_jit.py; optional host.py launch tuning." \
  --test-command "python3 scripts/task_runner.py correctness && python3 scripts/task_runner.py performance" \
  --mode full --gpu-ids 0 --no-rag

Full PerfSkills brief: ORIGINAL_TASK.md.