[albertvillanova]

This PR makes trl.kernels.selective_log_softmax_and_entropy compile with torch.compile(fullgraph=True) on torch 2.6.

Fix #7715.

Motivation

On torch 2.6, Dynamo can't trace Tensor.stride() inside the backward of a torch.autograd.Function in fullgraph mode. The backward of _LogProbEntropyFunction read the strides of index, row_mask and the gradient tensor, so compiling with fullgraph=True raised Unsupported: Illegal getattr invocation stride in strict mode. This was found while working on #7642 (torch 2.6.0 floor).

Solution

The backward no longer calls .stride(). It reuses the strides computed in the forward pass, and since the gradient tensor is allocated contiguous, its strides come from its shape.

Changes

• Save the index and row_mask strides on ctx in the forward pass and reuse them in the backward pass.
• Derive the strides of the contiguous grad_logits from its shape.

---

Note

Low Risk
Stride plumbing only; kernel math and pointer logic are unchanged, with contiguous grad_logits strides matching the prior _layout result.

Overview
Fixes torch.compile(fullgraph=True) failures for selective_log_softmax_and_entropy on PyTorch 2.6, where Dynamo rejects .stride() inside the backward of _LogProbEntropyFunction.

Forward now stores index and mask batch/row strides on ctx (index_strides, mask_strides) alongside the existing layout metadata. Backward reads those saved values instead of recomputing strides via _layout() (which called tensor.stride()). For the contiguous grad_logits buffer, strides are derived from shape as (rows_per_batch * n_cols, n_cols) rather than introspecting the tensor. Mask pointer selection is unchanged in behavior (index when mask is absent, else row_mask).

^{Reviewed by Cursor Bugbot for commit 9c8a235. Bugbot is set up for automated code reviews on this repo. Configure here.}

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[kashif]

checked also on ROCM