[FEAT][kernels]: add backward support for fused logp CUDA kernels

[Dnoob]

Closes #100

Summary

• Add CUDA backward kernels for fused logp, including indexed variants.
• Route the generic two-pass, online, and indexed allocating APIs through torch.autograd.Function when gradients are required.
• Make the experimental SM90 TMA path differentiable by returning softmax statistics from its forward pass and reusing the generic CUDA backward kernel.
• Preserve the existing no-grad forward results and output dtype behavior.

Implementation details

• Save row_max and log_sum separately to reconstruct probabilities without losing precision under large constant logit shifts.
• Compute grad_out * (one_hot(token_id) - softmax(logits)) without materializing a logits-sized probability tensor.
• Return exactly zero gradients for rows omitted by indexed variants.
• Keep caller-provided *_out and deterministic variants forward-only.
• Add extension bindings, type stubs, documentation, and focused gradient and regression tests.

SM90 fixes

H100 validation also exposed issues in the existing experimental SM90 forward path:

• Limit the TMA tile width to the supported CUtensorMap box dimension.
• Account for the full configured tile in the TMA transaction barrier, including a partial last tile.
• Replace the block-wide reduction with consumer-only warp reductions.
• Fall back to the generic CUDA backend when the bf16 row stride is not 16-byte aligned.

Validation

The backward test suite was run against the compiled CUDA extension on an NVIDIA H100.

root@b34a45b9a263:/workspace/RL-Kernel# nvidia-smi
Sun Jul 19 09:46:37 2026
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 580.126.09             Driver Version: 580.126.09     CUDA Version: 13.0     |
+-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|=========================================+========================+======================|
|   0  NVIDIA H100 80GB HBM3          On  |   00000000:AB:00.0 Off |                    0 |
| N/A   27C    P0             69W /  700W |       0MiB /  81559MiB |      0%      Default |
+-----------------------------------------+------------------------+----------------------+

root@b34a45b9a263:/workspace/RL-Kernel# python -c "import torch; print(f'PyTorch {torch.__version__}, CUDA {torch.version.cuda}, capability {torch.cuda.get_device_capability()}')"
PyTorch 2.4.1+cu124, CUDA 12.4, capability (9, 0)

root@b34a45b9a263:/workspace/RL-Kernel# PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 RL_KERNEL_REQUIRE_EXT=1 CUDA_LAUNCH_BLOCKING=1 python -m pytest tests/test_fused_logp_backward.py
================================================= test session starts ==================================================
platform linux -- Python 3.11.10, pytest-9.1.1, pluggy-1.6.0
rootdir: /workspace/RL-Kernel
configfile: pyproject.toml
collected 45 items

tests/test_fused_logp_backward.py ............................................. [100%]

================================================== 45 passed in 2.52s ==================================================

Summary by CodeRabbit

• New Features
  • Added LSE-enabled fused log-prob forward APIs that optionally return per-row row_max and log_sum (SM90/Hopper, indexed, and online variants).
  • Added fused log-prob backward APIs that use these saved LSE components for more efficient gradient computation.
• Documentation
  • Updated the fused log-prob operator docs to clarify tensor contracts, backend selection, and autograd/differentiability behavior (including forward-only *_out variants).
• Tests
  • Added a CUDA-focused backward test suite validating forward/LSE correctness, gradient accuracy, invalid-target handling, and SM90 routing/fallback.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: d98040c4-46dc-4d6d-8b6f-0c90b3ec231f

📥 Commits

Reviewing files that changed from the base of the PR and between 6536aed and 97ef7c6.

📒 Files selected for processing (2)

• csrc/ops.cpp
• rl_engine/_C.pyi

🚧 Files skipped from review as they are similar to previous changes (2)

• csrc/ops.cpp
• rl_engine/_C.pyi

---

📝 Walkthrough

Walkthrough

Fused LogP now exposes per-row log-sum-exp statistics and supports differentiable generic CUDA and SM90 execution. New backward kernels consume saved statistics, bindings and type stubs expose the APIs, Python routing selects autograd paths, and CUDA tests validate forward, backward, indexed, and SM90 behavior.

Changes

Fused LogP LSE and autograd

Layer / File(s)	Summary
Generic LSE and backward kernels csrc/fused_logp_kernel.cu	Forward variants optionally return row maxima and log-sums; regular and indexed backward kernels compute gradients from those saved statistics.
SM90 tiled LSE path csrc/cuda/fused_logp_sm90.cu	SM90 reductions and tile synchronization are revised, with a new LSE-returning launcher and forward API.
Extension and autograd wiring csrc/ops.cpp, rl_engine/_C.pyi, rl_engine/kernels/ops/cuda/loss/logp.py	New extension symbols and stubs are registered, and generic/SM90 operations select differentiable paths when gradients are requested.
Contracts and validation coverage docs/operators/fused-logp.md, tests/test_fused_logp_backward.py	Documentation describes tensor and gradient contracts; tests cover LSE outputs, backward variants, routing, invalid targets, indexed rows, and SM90 cases.

Estimated code review effort: 5 (Critical) | ~120 minutes

Sequence Diagram(s)

sequenceDiagram
  participant LogpOp
  participant CUDAForward
  participant CUDABackward
  LogpOp->>CUDAForward: request logp and LSE statistics
  CUDAForward-->>LogpOp: return output, row_max, log_sum
  LogpOp->>CUDABackward: pass gradients and saved statistics
  CUDABackward-->>LogpOp: return logits gradients

Loading

Possibly related issues

• [WS1 Test] #154 dtype pin + #152 KV path + #153 bwd invariance + #150 chain gate #244 — Adds and tests fused log-probability backward support and related gradient behavior.
• [Roadmap] RL-Kernel Roadmap Q3 - Q4 #83 — Tracks the fused LogP backward and exported-LSE roadmap implemented by this change.

Possibly related PRs

• RL-Align/RL-Kernel#122 — Updates the same SM90 fused LogP implementation and overlaps with this PR’s SM90 path changes.

Suggested reviewers: inaniloquentee, maxiaosong1124, flink-ddd, kjldefeated, zhangj1an

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely summarizes the main change: adding backward support for fused logp CUDA kernels.
Linked Issues check	✅ Passed	The PR implements the requested backward functions, autograd integration, bindings, and gradient tests for generic and SM90 kernels.
Out of Scope Changes check	✅ Passed	The added docs, stubs, LSE-returning helpers, and test coverage all support the backward-feature work and do not appear unrelated.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

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 `@csrc/cuda/fused_logp_sm90.cu`:
- Around line 142-144: Guard the selected-token lookup in the SM90 kernel around
label_idx so labels below zero or at least vocab_size produce the generic
backend’s 0 result without dereferencing logits_gmem. Also promote the row_idx *
vocab_size address calculation to a sufficiently wide index type before adding
label_idx, while preserving the existing valid-label log-probability
computation.

🪄 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: f2d7dd4c-67e2-4912-b97b-c724e0cd388c

📥 Commits

Reviewing files that changed from the base of the PR and between 6df029a and 5d86af8.

📒 Files selected for processing (7)

• csrc/cuda/fused_logp_sm90.cu
• csrc/fused_logp_kernel.cu
• csrc/ops.cpp
• docs/operators/fused-logp.md
• rl_engine/_C.pyi
• rl_engine/kernels/ops/cuda/loss/logp.py
• tests/test_fused_logp_backward.py