Skip to content

Move measurement out of kernel.py for remaining GEAK tasks#36

Merged
irvineoy merged 1 commit into
AMD-AGI:mainfrom
sdubagun-amd:standardize-exploit-relevant-harnesses
Jul 14, 2026
Merged

Move measurement out of kernel.py for remaining GEAK tasks#36
irvineoy merged 1 commit into
AMD-AGI:mainfrom
sdubagun-amd:standardize-exploit-relevant-harnesses

Conversation

@sdubagun-amd

@sdubagun-amd sdubagun-amd commented May 6, 2026

Copy link
Copy Markdown
Collaborator

Moves correctness and measurement ownership out of the agent-editable kernel.py for the two remaining tasks on current main:

  • lean_atten_paged
  • gemm_a16wfp4

Changes:

  • lean_atten_paged: remove correctness, profiling, benchmark, speedup, and CLI code from kernel.py; keep the PyTorch reference and CUDA Event median timing in test_kernel_harness.py.
  • gemm_a16wfp4: replace host perf_counter timing with CUDA Event median timing in test_kernel_harness.py.
  • Preserve the current main benchmark shape sets and emit canonical ALL_SHAPES indices through GEAK_SHAPES_USED.
  • Keep the gemm FP4 capability check in the harness rather than trusting the agent-editable kernel module.
  • Emit only GEAK_RESULT_LATENCY_MS; baseline-versus-patched speedup remains the orchestrator responsibility.

The earlier refk_identity and refk_fp8_blockwise_mm changes are omitted because equivalent harness-integrity fixes already landed on main. ff_backward is omitted because the task was removed from main.

Validation:

  • MI355X/gfx950 with lmsysorg/sglang-rocm:v0.5.14-rocm720-mi35x-20260705.
  • lean_atten_paged correctness: 4/4 passed.
  • gemm_a16wfp4 correctness: 55/55 passed.
  • Full 50-warmup/200-iteration benchmark passed on all 7 lean_atten_paged shapes; geometric mean latency was 0.0883 ms.
  • Full 50-warmup/200-iteration benchmark passed on all 55 gemm_a16wfp4 shapes; geometric mean latency was 0.1696 ms.
  • Both benchmarks emitted canonical GEAK_SHAPES_USED values and ended with GEAK_RESULT_LATENCY_MS.
  • Profile smoke tests passed on 5/5 configs for each task with 1 warmup and 1 measured iteration.
  • Python AST parsing, harness ownership assertions, and Docker CLI/import checks passed.

The latency values above are functional benchmark validation results, not controlled leaderboard measurements; GPU clocks were not locked.

@irvineoy

irvineoy commented May 6, 2026

Copy link
Copy Markdown
Collaborator

One issue I noticed with the sampled harnesses: GEAK_SHAPES_USED no longer reports the original ALL_SHAPES indices for tasks that downsample.

For example:

  • refk_fp8_blockwise_mm/test_kernel_harness.py builds HARNESS_SHAPES from _harness_indices, but later prints GEAK_SHAPES_USED={list(range(len(shapes)))}.
  • gemm_a16wfp4/test_kernel_harness.py does the same.

On this PR, the actual sampled source indices are different from the emitted values:

  • refk_fp8_blockwise_mm: sampled [0, 1, 2, 4, 5, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 18, 19, 20, 21, 22, 23, 24, 26, 27, 28], but reports [0..24].
  • gemm_a16wfp4: sampled [0, 2, 4, 7, 9, 11, 14, 16, 18, 20, 22, 25, 27, 29, 32, 34, 36, 38, 40, 43, 45, 47, 50, 52, 54], but reports [0..24].

If the orchestrator trusts GEAK_SHAPES_USED for attribution or replaying base-vs-patched comparisons, this can point it at the wrong shapes. Could we preserve/pass the selected source indices and print those instead of the local positions?

Keep correctness and timing in the harness for lean_atten_paged and standardize gemm_a16wfp4 on CUDA event median timing. Preserve the current main shape sets, report canonical shape indices, and keep the FP4 capability check outside the agent-editable kernel module.

Co-authored-by: Cursor <cursoragent@cursor.com>
@irvineoy
irvineoy force-pushed the standardize-exploit-relevant-harnesses branch from feb30e6 to 2ea64ed Compare July 14, 2026 03:34
@irvineoy irvineoy changed the title Move measurement out of kernel.py into the harness for 5 tasks Move measurement out of kernel.py for remaining GEAK tasks Jul 14, 2026
@irvineoy
irvineoy changed the base branch from geak-triton-common-benchmark to main July 14, 2026 03:35
@irvineoy
irvineoy merged commit faf6564 into AMD-AGI:main Jul 14, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants