nll_loss backward: store address built from the ignore_index sentinel can alias another row's gradient

[truong-v]

Describe the bug

nll_loss_backward_kernel and nll_loss2d_backward_kernel in src/flag_gems/ops/nllloss.py gate the gradient value on ignore_mask, but not the address they write it to, nor the store's mask:

tgt = tl.load(tgt_ptr + offsets_n, mask=mask_n, other=0)
ignore_mask = not (tgt == ignore_index) and mask_n
...
inp_grad = tl.where(ignore_mask, -1 * out_grad * wgt_tgt / total_w, 0)   # value: gated
inp_grad_ptrs = inp_grad_ptr + offsets_n * C + tgt                        # address: NOT gated
tl.store(inp_grad_ptrs, inp_grad, mask=mask_n)                            # mask: row-bound only

On an ignored row, tgt holds ignore_index itself. The file's own comment (line ~237) states the contract:

If ignore_index is specified, this index can be outside the class range and will be ignored in the loss computation.

So tgt can legitimately be outside [0, C) — torch.nn.functional.nll_loss's default is -100, and writing that sentinel into target is the standard way to mark padded or masked positions. The kernel nevertheless builds offsets_n * C + tgt from it and stores unconditionally under mask_n, which only asks "is this row within N", never "is this row ignored".

Two consequences:

1. A positive out-of-range sentinel aliases another row's gradient cell. With C = 4 and ignore_index = 4, ignored row n addresses n*4 + 4, which is row n+1's cell 0. Two threads then write different values to that one address in the same instruction, and which one survives is whatever the hardware schedules last. Observed result: the valid row's real gradient is replaced by the ignored row's 0.
2. A negative sentinel writes outside grad_input. ignore_index = -100 addresses 100 elements before the tensor. In my runs this landed in allocator slack and produced correct-looking output, but it is still an out-of-bounds write.

nll_loss_nd_backward_kernel in ops/nll_loss_nd.py does not have this problem — it masks its store on target validity (mask=valid).

Steps to reproduce

import torch
import flag_gems

N, C = 6, 4
ignore_index = C  # outside [0, C), which the op's own contract allows

inp = torch.randn(N, C, dtype=torch.float32, device="cuda", requires_grad=True)
target = torch.tensor([ignore_index, 0, 1, ignore_index, 0, 3], device="cuda")
out_grad = torch.ones(N, device="cuda")

ref = torch.nn.functional.nll_loss(inp, target, reduction="none",
                                   ignore_index=ignore_index)
(ref_grad,) = torch.autograd.grad(ref, inp, out_grad, retain_graph=True)

with flag_gems.use_gems():
    res = torch.nn.functional.nll_loss(inp, target, reduction="none",
                                       ignore_index=ignore_index)
    (res_grad,) = torch.autograd.grad(res, inp, out_grad)

print("torch:\n", ref_grad)
print("gems:\n", res_grad)
print("max abs diff:", (res_grad - ref_grad).abs().max().item())

Output on bcf3f5b7:

torch:
 tensor([[ 0.,  0.,  0.,  0.],
        [-1.,  0.,  0.,  0.],
        [ 0., -1.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [-1.,  0.,  0.,  0.],
        [ 0.,  0.,  0., -1.]], device='cuda:0')
gems:
 tensor([[ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0., -1.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0., -1.]], device='cuda:0')
max abs diff: 1.0

Rows 1 and 4 are ordinary valid rows, and their gradients are gone: ignored row 0 wrote a 0 over row 1's cell 0, and ignored row 3 over row 4's cell 0. Rows 2 and 5 are untouched because no ignored row happens to alias them.

The same shape of hazard is in nll_loss2d_backward_kernel (line ~220, offset_n * C * D + tgt * D + offset_d, stored under mask_block): an ignored pixel of sample n addresses sample n+1's class-0 gradient at the same pixel.

Expected behavior

An ignored row contributes nothing to grad_input, and cannot affect any other row's gradient. grad_input is torch.zeros_like-allocated, so an ignored row needs no store at all.

Why the current tests do not catch it

All five nll_loss test files build targets with target = torch.randint(0, shape[dim], target_shape, ...), i.e. always inside [0, C). The sentinel therefore never appears in target and never reaches the address arithmetic:

• ignore_index=-100 (and 200 for C=2) never matches any target, so no row is ever ignored.
• ignore_index=1 (and 200 for C=256) does mark rows ignored, but it is a valid class id, so n*C + 1 is that row's own cell and writing 0 there is correct.

The full tests/test_nll_loss_backward.py suite passes on unfixed code for this reason.

Suggested fix

Keep the sentinel out of the address arithmetic and predicate the store on ignore_mask:

safe_tgt = tl.where(ignore_mask, tgt, 0)
inp_grad_ptrs = inp_grad_ptr + offsets_n * C + safe_tgt
tl.store(inp_grad_ptrs, inp_grad, mask=ignore_mask)

and the same for nll_loss2d_backward_kernel's offset_n * C * D + tgt * D + offset_d. This matches what nll_loss_nd_backward_kernel already does.

I have a patch with regression tests for both kernels and will open a PR referencing this issue.

Environment

• FlagGems commit bcf3f5b7
• torch 2.6.0+cu124, triton 3.2.0
• NVIDIA H100 NVL