Add GRPO loss example

[ighoshsubho]

Overview

This PR adds a complete example implementation of Group Relative Policy Optimization (GRPO) loss using Helion kernels, including:

• Helion kernels for forward and backward passes
• PyTorch reference implementation for validation
• Autograd integration via a custom torch.autograd.Function
• Unit-style verification and a micro-benchmark harness

Benchmarks

=== Timing (median ms) ===
PyTorch Forward: 5.286 ms
PyTorch Backward: 14.828 ms
Helion Forward: 1.053 ms (x5.02 vs Torch)
Helion Backward: 2.461 ms (x6.03 vs Torch)

=== Throughput ===
PyTorch Fwd tokens/s: 3099305.7
PyTorch Bwd tokens/s: 1104966.6
Helion Fwd tokens/s: 15564766.3
Helion Bwd tokens/s: 6658797.2

Motivation

GRPO is used to optimize RLHF-style training by stabilizing policy updates through clipping with optional KL regularization against a reference model. This example demonstrates how to express a numerically stable, fused GRPO loss in Helion, and how to integrate Helion kernels into an end-to-end PyTorch workflow.

[jansel]

Thanks for the contribution! Can you add a test for this similar to the tests for the other examples?

[ighoshsubho]

Thanks for the contribution! Can you add a test for this similar to the tests for the other examples?

Yeah @jansel I added them just now with changes in both .expected file and the examples!

[ighoshsubho]

@oulgen I have removed the redundant configs for a specific hardware. Is it ok to merge?

[jansel]

Looks like many of the tests are failing with:
FAILED test/test_examples.py::TestExamples::test_grpo_loss_bwd - Failed: Timeout (>60.0s) from pytest-timeout.

Can we change the test to run faster? Perhaps smaller inputs?

[ighoshsubho]

I have made the tensor shape small enough to fit under 60 sec runtime compilation, should not throw any error for backward now.

cc: @jansel @oulgen

[oulgen]

I have made the tensor shape small enough to fit under 60 sec runtime compilation, should not throw any error for backward now.

cc: @jansel @oulgen

thanks, i'll merge if/when the tests pass