[FEAT][kernels]: implement Fused Policy Ratio and KL Penalty kernel - #97
Conversation
…Triton/CUDA/ROCm)
📝 WalkthroughWalkthroughThis PR introduces a fused ChangesFused Ratio-KL and GRPO Loss Integration
Sequence Diagram(s)sequenceDiagram
participant User
participant NativeGRPOLossOp
participant NativeRatioKLOp
participant Advantage as Advantage Normalization
participant Loss as Loss Computation
User->>NativeGRPOLossOp: forward(policy_logits, ref_logits, action_ids, old_logps, rewards)
NativeGRPOLossOp->>Advantage: group_advantages(rewards, boundaries)
Advantage-->>NativeGRPOLossOp: sample_advantages
NativeGRPOLossOp->>NativeRatioKLOp: __call__(policy_logits, ref_logits, action_ids, ...)
NativeRatioKLOp->>NativeRatioKLOp: select log-probs, compute delta/diff
NativeRatioKLOp-->>NativeGRPOLossOp: ratio, kl_terms
NativeGRPOLossOp->>Loss: expand_advantages, compute clipped surrogate
Loss-->>NativeGRPOLossOp: loss, policy_loss, kl
NativeGRPOLossOp-->>User: (loss, policy_loss, kl), backward graph
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related issues
Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rl_engine/kernels/ops/pytorch/loss/grpo_loss.py (1)
192-203:⚠️ Potential issue | 🟠 Major | ⚡ Quick winMirror Triton’s
group_boundariesvalidation here.
_resolve_group_ids()currently accepts translated boundary vectors because it only checkssizes. For example, a 4-sequence batch would accept[2, 4, 6]and silently treat it like[0, 2, 4], whileTritonGRPOLossOp._build_bounds()rejects the same input because it requires batch-local offsets that start at0and end atnum_sequences. That makes the shared API behave differently across backends depending on dispatch target.Suggested fix
boundaries = torch.as_tensor(group_boundaries, device=device, dtype=torch.long) if boundaries.ndim != 1 or boundaries.numel() < 2: raise ValueError("group_boundaries must be a 1D tensor of length num_groups + 1.") + if int(boundaries[0].item()) != 0 or int(boundaries[-1].item()) != num_sequences: + raise ValueError("group_boundaries must start at 0 and end at num_sequences.") sizes = boundaries[1:] - boundaries[:-1] - - if int(sizes.sum().item()) != num_sequences: - raise ValueError( - f"group sizes sum to {int(sizes.sum().item())} but there are " - f"{num_sequences} sequences." - ) if bool((sizes < 1).any().item()): raise ValueError("each group must contain at least one sequence.")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rl_engine/kernels/ops/pytorch/loss/grpo_loss.py` around lines 192 - 203, _update _resolve_group_ids() to mirror TritonGRPOLossOp._build_bounds() validation: ensure group_boundaries is a 1D integer tensor of length num_groups+1 on the correct device, that boundaries[0] == 0 and boundaries[-1] == num_sequences, and that the sequence is non-decreasing (each boundary >= previous) and all group sizes are >=1; raise a ValueError with a clear message if any of these checks fail so the PyTorch path rejects translated/offset vectors like [2,4,6] just like Triton does.
🧹 Nitpick comments (2)
docs/operators/ratio-kl.md (1)
9-11: 💤 Low valueConsider adding a language identifier to the fenced code block.
The ASCII pipeline diagram lacks a language specifier. Adding
textimproves rendering and addresses the linter warning.📝 Proposed fix
-``` +```text logits --[ratio_kl op]--> (ratio, kl) --[clipped surrogate + beta*kl]--> loss</details> <details> <summary>🤖 Prompt for AI Agents</summary>Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.In
@docs/operators/ratio-kl.mdaround lines 9 - 11, The fenced ASCII pipeline
block containing "logits --[ratio_kl op]--> (ratio, kl) --[clipped surrogate +
beta*kl]--> loss" should include a language identifier (e.g., text) to satisfy
the linter and improve rendering; update the triple backtick fence that
surrounds the pipeline diagram to use ```text so the block becomes a text code
block while keeping the pipeline content unchanged.</details> <!-- cr-comment:v1:ade0dc832571fe86b6ad1795 --> _Source: Linters/SAST tools_ </blockquote></details> <details> <summary>docs/operators/grpo-loss.md (1)</summary><blockquote> `15-15`: _💤 Low value_ **Consider adding a language identifier to the fenced code block.** The ASCII pipeline diagram uses a fenced code block without a language specifier. While this is valid Markdown, adding `text` improves rendering consistency and silences the linter warning. <details> <summary>📝 Proposed fix</summary> ```diff -``` +```text logits --[ratio_kl op]--> (ratio, kl) --[group adv + clipped surrogate]--> loss ``` ``` </details> <details> <summary>🤖 Prompt for AI Agents</summary>Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.In
@docs/operators/grpo-loss.mdat line 15, The fenced ASCII pipeline diagram
lacks a language tag; update the fenced code block that contains "logits
--[ratio_kl op]--> (ratio, kl) --[group adv + clipped surrogate]--> loss" to
include a language identifier (use "text") after the opening backticks so the
block becomestext ...to satisfy the linter and improve rendering
consistency.</details> <!-- cr-comment:v1:58983bcca01d03c1472ca2d1 --> _Source: Linters/SAST tools_ </blockquote></details> </blockquote></details> <details> <summary>🤖 Prompt for all review comments with AI agents</summary>Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.Inline comments:
In@benchmarks/benchmark_grpo_loss.py:
- Around line 93-109: The nested functions native_fwd, triton_fwd,
native_fwd_bwd, and triton_fwd_bwd capture loop-scoped variables call_args and
spp causing potential late-binding (Ruff B023); fix by binding those loop
variables as default arguments in each function signature (e.g.,
native_fwd(p=policy, call_args=call_args, spp=spp) and similarly triton_fwd, and
for the backward helpers use p=pol_grad, call_args=call_args, spp=spp) so each
closure retains the current iteration's values instead of the final ones.In
@docs/operators/grpo-loss.md:
- Line 30: Doc parameter names are inconsistent: grpo_loss docs use
completion_mask while ratio_kl docs and implementations use attention_mask;
update the docs so the same tensor name is used and unambiguous. Pick one
canonical name (e.g., attention_mask) and rename the parameter in
docs/operators/grpo-loss.md (and any related doc examples) from completion_mask
to attention_mask, or add a single clarifying sentence in grpo_loss docs stating
that completion_mask is the same tensor passed as attention_mask to ratio_kl;
ensure references to shapes ([B, T]) and callsites in grpo_loss, ratio_kl, and
fused _ratio_kl are updated to match the chosen terminology.In
@docs/operators/ratio-kl.md:
- Around line 43-48: Add a short introductory sentence before the formula to
explain what it represents (for example: "The Triton backward kernel
computes:"), add a language identifier to the fenced code block (e.g.,text orpython) around the linegrad_policy_logits[v] = c * (1[v == action] - softmax_policy(v)), and either move or expand the following sentence ("so the
backward also avoids materializing...") so it immediately relates to the formula
(e.g., state that this expression means no [B, T, V] probability tensor is
materialized), ensuring the formula is connected to the preceding discussion of
backends and wrapper classes.In
@rl_engine/kernels/ops/triton/loss/ratio_kl.py:
- Around line 38-67: The kernel reads an unclamped action id "a" (loaded as a =
tl.load(action_ptr + row)) and then uses it to index policy_ptr/ref_ptr, which
can cause out-of-bounds reads for active tokens; clamp "a" to the valid range
[0, V-1] before any indexing (e.g., replace uses of "a" with a_safe =
tl.minimum(tl.maximum(a, 0), V - 1) or tl.clamp if available) and ensure a_safe
has the correct integer dtype matching row_off offsets; update the loads
tl.load(policy_ptr + row_off + a) and tl.load(ref_ptr + row_off + a) to use
a_safe so active-path indexing is always safe in the ratio_kl Triton forward
kernel.
Outside diff comments:
In@rl_engine/kernels/ops/pytorch/loss/grpo_loss.py:
- Around line 192-203: _update _resolve_group_ids() to mirror
TritonGRPOLossOp._build_bounds() validation: ensure group_boundaries is a 1D
integer tensor of length num_groups+1 on the correct device, that boundaries[0]
== 0 and boundaries[-1] == num_sequences, and that the sequence is
non-decreasing (each boundary >= previous) and all group sizes are >=1; raise a
ValueError with a clear message if any of these checks fail so the PyTorch path
rejects translated/offset vectors like [2,4,6] just like Triton does.
Nitpick comments:
In@docs/operators/grpo-loss.md:
- Line 15: The fenced ASCII pipeline diagram lacks a language tag; update the
fenced code block that contains "logits --[ratio_kl op]--> (ratio, kl) --[group
adv + clipped surrogate]--> loss" to include a language identifier (use "text")
after the opening backticks so the block becomestext ...to satisfy the
linter and improve rendering consistency.In
@docs/operators/ratio-kl.md:
- Around line 9-11: The fenced ASCII pipeline block containing "logits
--[ratio_kl op]--> (ratio, kl) --[clipped surrogate + beta*kl]--> loss" should
include a language identifier (e.g., text) to satisfy the linter and improve
rendering; update the triple backtick fence that surrounds the pipeline diagram
to use ```text so the block becomes a text code block while keeping the pipeline
content unchanged.</details> <details> <summary>🪄 Autofix (Beta)</summary> Fix all unresolved CodeRabbit comments on this PR: - [ ] <!-- {"checkboxId": "4b0d0e0a-96d7-4f10-b296-3a18ea78f0b9"} --> Push a commit to this branch (recommended) - [ ] <!-- {"checkboxId": "ff5b1114-7d8c-49e6-8ac1-43f82af23a33"} --> Create a new PR with the fixes </details> --- <details> <summary>ℹ️ Review info</summary> <details> <summary>⚙️ Run configuration</summary> **Configuration used**: defaults **Review profile**: CHILL **Plan**: Pro **Run ID**: `b91d50a6-caba-4711-b397-1074cc35634b` </details> <details> <summary>📥 Commits</summary> Reviewing files that changed from the base of the PR and between 4a9ca42a07c9da6b5634375e1043f9150b19a9e7 and 6a858ae077e3302f6c0456911a289f4cb11481e5. </details> <details> <summary>📒 Files selected for processing (14)</summary> * `benchmarks/benchmark_grpo_loss.py` * `benchmarks/benchmark_ratio_kl.py` * `docs/.nav.yml` * `docs/operators/README.md` * `docs/operators/grpo-loss.md` * `docs/operators/ratio-kl.md` * `rl_engine/kernels/ops/pytorch/loss/grpo_loss.py` * `rl_engine/kernels/ops/pytorch/loss/ratio_kl.py` * `rl_engine/kernels/ops/triton/loss/__init__.py` * `rl_engine/kernels/ops/triton/loss/grpo_loss.py` * `rl_engine/kernels/ops/triton/loss/ratio_kl.py` * `rl_engine/kernels/registry.py` * `tests/test_grpo_loss.py` * `tests/test_ratio_kl.py` </details> </details> <!-- This is an auto-generated comment by CodeRabbit for review status -->
Flink-ddd
left a comment
There was a problem hiding this comment.
Overall clean, online-softmax fwd and the backward chain rule look correct, nice to see the bespoke GRPO kernels removed. One blocking item (action_ids clamp parity with the native op) plus two test-coverage suggestions inline. Happy to approve once the clamp is addressed.
| pol = policy_logits.contiguous().view(-1, V) | ||
| ref = ref_logits.contiguous().view(-1, V) | ||
| n_rows = pol.shape[0] | ||
| act = action_ids.contiguous().view(-1).to(torch.int64) |
There was a problem hiding this comment.
[blocking] action_ids not clamped before indexing, diverges from native op
The native op clamps action ids before gathering:
safe_ids = action_ids.clamp(0, logits.size(-1) - 1).long()
but the Triton path indexes policy_ptr + row_off + a directly. For an active row with an out-of-range action_id (corrupt/unexpected data), this is an illegal memory access rather than a clamped read — the two backends behave differently on the same input. Suggest clamping host-side to keep parity:
act = action_ids.contiguous().view(-1).clamp_(0, V - 1).to(torch.int64)
| assert pol_t.grad is not None | ||
| assert torch.isfinite(pol_t.grad).all() | ||
| assert torch.allclose(pol_t.grad, pol_n.grad, atol=1e-4, rtol=1e-4) | ||
|
|
There was a problem hiding this comment.
[test gap] No assertion that ref_logits receives no gradient on the Triton path
test_native_gradient_flows_to_policy_logits asserts ref_logits.grad is None, but the Triton backward tests never set requires_grad_ on ref_logits, so the freeze isn't actually verified for the fused op. Suggest a symmetric test:
@requires_triton_cuda
def test_triton_no_grad_to_ref():
fused = TritonRatioKLOp()
pol, ref, act, mask, old = _inputs(seed=8, device="cuda")
pol = pol.clone().requires_grad_(True)
ref = ref.clone().requires_grad_(True)
r, k = fused(pol, ref, act, mask, old)
(r.sum() + k.sum()).backward()
assert ref.grad is None
| assert torch.allclose(base_k[active], pert_k[active], atol=1e-5) | ||
| assert torch.allclose(pert_r[inactive], torch.ones_like(pert_r[inactive])) | ||
|
|
||
|
|
There was a problem hiding this comment.
[test gap] No coverage for out-of-range action_ids / clamp behavior
Related to the clamp divergence above. Once the behavior is settled (clamp on both paths), add a test that puts an out-of-range id at a masked position and asserts both backends stay finite and agree on active outputs:
@requires_triton_cuda
def test_triton_handles_oob_action_ids():
native = NativeRatioKLOp()
fused = TritonRatioKLOp()
pol, ref, act, mask, old = _inputs(seed=9, device="cuda", valid_density=0.7)
act = act.clone()
inactive = ~mask.to(torch.bool)
act[inactive] = _VOCAB + 999 # garbage id at masked positions
r_t, k_t = fused(pol, ref, act, mask, old)
r_n, k_n = native(pol, ref, act, mask, old)
active = mask.to(torch.bool)
assert torch.isfinite(r_t).all() and torch.isfinite(k_t).all()
assert torch.allclose(r_t[active], r_n[active], atol=1e-4, rtol=1e-4)
assert torch.allclose(k_t[active], k_n[active], atol=1e-4, rtol=1e-4)
…RL-Kernel into feat/fused-ratio-kl
|
@Flink-ddd The requests have been resolved. Thx for reviews. |
Flink-ddd
left a comment
There was a problem hiding this comment.
LGTM now, but also need @inaniloquentee to review.
|
The GPU CI test is unrelated to the code because GPU resources are temporarily occupied, we'll try again later. |
|
|
||
| class _RatioKLFunction(torch.autograd.Function): | ||
| @staticmethod | ||
| def forward(ctx, policy_logits, ref_logits, action_ids, attention_mask, old_logps): |
There was a problem hiding this comment.
Could we add a quick shape check before flattening here? If action_ids, attention_mask, or old_logps don't match policy_logits.shape[:-1], the Triton kernel can still launch and read the wrong rows.
There was a problem hiding this comment.
Thx for suggestion, I will fix the requests.
|
|
||
| def _selected_logp(self, logits: torch.Tensor, action_ids: torch.Tensor) -> torch.Tensor: | ||
| # Clamp ids first so masked/pad positions never index out of bounds | ||
| safe_ids = action_ids.clamp(0, logits.size(-1) - 1).long() |
There was a problem hiding this comment.
Non-blocking, but could we avoid clamping active token ids here? Clamping is fine for masked positions, but an out-of-range active action_id may be better surfaced as an error instead of silently using the last vocab entry.
There was a problem hiding this comment.
I see, I will use a try-catch to deal with it
|
@inaniloquentee the requests are resolved, thx for reviews! |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rl_engine/kernels/ops/pytorch/loss/ratio_kl.py`:
- Around line 41-50: The validation check for action_ids uses the mask to only
validate active (unmasked) positions, allowing out-of-range IDs at masked
positions to pass through and cause undefined behavior in torch.gather(). Remove
the mask from the validation logic in the line computing the invalid variable so
that ALL action_ids are checked to be within the valid range [0, vocab_size),
regardless of masking. Additionally, validate that both policy_logits and
ref_logits have the same vocabulary size before using them in the
_selected_logp() method calls, ensuring consistency when indexing both logit
tensors with the same action_ids.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 5cd66bfc-fca8-4da6-9673-341eb98b4f57
📒 Files selected for processing (2)
rl_engine/kernels/ops/pytorch/loss/ratio_kl.pyrl_engine/kernels/ops/triton/loss/ratio_kl.py
🚧 Files skipped from review as they are similar to previous changes (1)
- rl_engine/kernels/ops/triton/loss/ratio_kl.py
| vocab_size = policy_logits.size(-1) | ||
| invalid = mask & ((action_ids < 0) | (action_ids >= vocab_size)) | ||
| if invalid.any(): | ||
| raise ValueError( | ||
| f"action_ids at active (unmasked) positions must be in " | ||
| f"[0, {vocab_size}); found out-of-range ids." | ||
| ) | ||
| logp_policy = self._selected_logp(policy_logits, action_ids) | ||
| with torch.no_grad(): | ||
| logp_ref = self._selected_logp(ref_logits, action_ids) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Verify whether NativeLogpOp.apply_fp32 performs raw gather/indexing or internal clamping.
fd -i 'logp.py' rl_engine/kernels/ops/pytorch/loss --exec sed -n '1,260p' {}
rg -n -C3 'def apply_fp32|gather|take_along_dim|index_select|clamp|action_ids' rl_engine/kernels/ops/pytorch/loss/logp.py
# Verify test coverage for masked out-of-range action_ids.
rg -n -C3 'out[- ]of[- ]range|masked|action_ids|mask' tests/test_ratio_kl.pyRepository: RL-Align/RL-Kernel
Length of output: 11647
🏁 Script executed:
python3 - <<'PY'
import torch
# Test 1: Does torch.gather fail on out-of-bounds indices?
try:
logits = torch.randn(2, 3, 100) # shape (batch=2, seq=3, vocab=100)
logits_2d = logits.reshape(-1, 100) # (6, 100)
# Valid indices
valid_ids = torch.tensor([5, 50, 99], dtype=torch.long)
result_valid = torch.gather(logits_2d, dim=-1, index=valid_ids.unsqueeze(1))
print("✓ Test 1a: Valid indices work:", result_valid.shape)
except Exception as e:
print(f"✗ Test 1a failed: {e}")
try:
# Out-of-bounds indices
oob_ids = torch.tensor([5, 150, 99], dtype=torch.long) # 150 >= vocab_size=100
result_oob = torch.gather(logits_2d, dim=-1, index=oob_ids.unsqueeze(1))
print("✓ Test 1b: OOB indices returned:", result_oob.shape)
print(" (This suggests torch.gather does NOT validate bounds at gather time)")
except IndexError as e:
print(f"✗ Test 1b: IndexError raised: {e}")
except Exception as e:
print(f"✗ Test 1b other error: {type(e).__name__}: {e}")
# Test 2: Check if result contains NaN or fails silently
print("\nTest 2: Behavior with mixed valid/OOB indices")
try:
logits = torch.randn(2, 1000)
ids = torch.tensor([100, 999, 5000], dtype=torch.long) # 5000 > vocab=1000
gathered = torch.gather(logits, dim=-1, index=ids.unsqueeze(1))
print(f" Gathered shape: {gathered.shape}")
print(f" Contains NaN: {torch.isnan(gathered).any().item()}")
print(f" Contains Inf: {torch.isinf(gathered).any().item()}")
print(f" Values: {gathered.squeeze()}")
except Exception as e:
print(f" Error: {type(e).__name__}: {e}")
PYRepository: RL-Align/RL-Kernel
Length of output: 180
🏁 Script executed:
# Check if the test actually runs and passes
cd tests && python3 -m pytest test_ratio_kl.py::test_triton_handles_oob_action_ids -v 2>&1 | head -50Repository: RL-Align/RL-Kernel
Length of output: 103
🏁 Script executed:
# Also verify the actual ratio_kl.py implementation - check if there's any clamping we missed
sed -n '1,60p' rl_engine/kernels/ops/pytorch/loss/ratio_kl.pyRepository: RL-Align/RL-Kernel
Length of output: 1972
Out-of-range action IDs at masked positions bypass validation and cause undefined behavior in torch.gather
Validation on line 41-47 only catches out-of-range IDs at active (unmasked) positions: invalid = mask & ((action_ids < 0) | (action_ids >= vocab_size)). IDs at masked positions are never validated. These unclamped IDs are then passed directly to torch.gather() on lines 48-50 before masking is applied, resulting in undefined behavior. Additionally, only policy_logits.size(-1) is validated, but ref_logits is indexed with the same action IDs.
Proposed fix
mask = attention_mask.to(torch.bool)
vocab_size = policy_logits.size(-1)
+ if ref_logits.size(-1) != vocab_size:
+ raise ValueError(
+ "policy_logits and ref_logits must have the same vocab dimension."
+ )
invalid = mask & ((action_ids < 0) | (action_ids >= vocab_size))
if invalid.any():
raise ValueError(
f"action_ids at active (unmasked) positions must be in "
f"[0, {vocab_size}); found out-of-range ids."
)
- logp_policy = self._selected_logp(policy_logits, action_ids)
+ safe_action_ids = torch.where(mask, action_ids, torch.zeros_like(action_ids))
+ logp_policy = self._selected_logp(policy_logits, safe_action_ids)
with torch.no_grad():
- logp_ref = self._selected_logp(ref_logits, action_ids)
+ logp_ref = self._selected_logp(ref_logits, safe_action_ids)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rl_engine/kernels/ops/pytorch/loss/ratio_kl.py` around lines 41 - 50, The
validation check for action_ids uses the mask to only validate active (unmasked)
positions, allowing out-of-range IDs at masked positions to pass through and
cause undefined behavior in torch.gather(). Remove the mask from the validation
logic in the line computing the invalid variable so that ALL action_ids are
checked to be within the valid range [0, vocab_size), regardless of masking.
Additionally, validate that both policy_logits and ref_logits have the same
vocabulary size before using them in the _selected_logp() method calls, ensuring
consistency when indexing both logit tensors with the same action_ids.
LGTM now. Thanks! |
#40
Summary
What changed and why
New operator —
ratio_klrl_engine/kernels/ops/triton/loss/ratio_kl.py:_RatioKLFunctionautograd +TritonRatioKLOp.BLOCK_Vtiles for policy and refrl_engine/kernels/ops/pytorch/loss/ratio_kl.py:NativeRatioKLOpreference, reusingNativeLogpOptests/test_ratio_kl.py: native-vs-reference, masked neutrality, ratio==1 whenold==policy, grad-to-policy-only, Triton fwd/bwd vs native (V=64 and 50257), registry dispatch.benchmarks/benchmark_ratio_kl.py: New benchmark (latency + peak VRAM, vocab sweep).docs/operators/ratio-kl.md: operator doc.GRPO loss refactor — compose
ratio_klinstead of a bespoke kernelrl_engine/kernels/ops/triton/{triton_grpo_loss.py → loss/grpo_loss.py}:loss/;_grpo_fwd/_grpo_bwdkernels and_GRPOLossFunction;TritonRatioKLOpand keeps only_group_norm_kernelfor reward normalization.benchmarks/benchmark_grpo_loss.py: Updated to logits APItests/test_grpo_loss.py: Loss/grad/masking/SGD tests reworked for logits APIdocs/operators/grpo-loss.mdWiring
rl_engine/kernels/registry.py: Registerratio_kl(Triton/PyTorch) for cuda/rocm/cpu; repointTRITON_GRPO_LOSStotriton.loss.grpo_lossPerformance (fp16, B=32, T=256)
Notes
Summary by CodeRabbit
Release Notes
New Features
ratio_kloperator for computing policy ratio and KL penalty in PPO/GRPO workflows.ratio_klpath.Documentation
ratio_kloperator documentation and linked it from the operators index.Tests
ratio_kltest suite (correctness, masking neutrality, gradient flow, and backend dispatch).Benchmarking
ratio_kl-driven composition and updated tensor shapes/CLI options.