[xpu-alignment] [26-11-11] XPU grid_sampler_2d backward accumulates the fp16/bf16 grid gradient in the input dtype: 2048 unit contributions sum to 1024 (fp16) and 128 (bf16) #6162

Summary

grid_sampler_2d_backward on XPU keeps its grid-gradient accumulator in the input dtype, so for float16 and bfloat16 the per-channel sum stops growing once each tap falls below half a unit in the last place of the running sum. With 2048 channels that each contribute exactly 1 to the x-gradient and a coordinate scale of 1, XPU returns [1024.0, 0.0] for float16 and [128.0, 0.0] for bfloat16, where the exact answer [2048.0, 0.0] is representable in both dtypes. float32 on the same device is exact, so this is accumulator loss, not representation error, and it is silent: no warning, no error.

The practical effect is that half-precision F.grid_sample backward reports grid gradients that are too small by up to 16x for wide channel counts, which quietly stalls learning of the sampling grid (spatial transformers, optical-flow warps, deformable sampling).

Upstream source

• [CUDA] grid_sampler_2d backward: FP16/BF16 grad_grid accumulates in the input dtype pytorch/pytorch#204496, "[CUDA] grid_sampler_2d backward: FP16/BF16 grad_grid accumulates in the input dtype", created 2026-11-11, open and labelled module: correctness (silent), module: numerical-stability, module: half, module: bfloat16, module: interpolation. The values were measured on ROCm (MI355X); the CUDA behaviour is derived from source and was not executed.
• [ROCm] grid_sampler_2d backward: accumulate FP16/BF16 grid gradient in FP32 (changes values) pytorch/pytorch#204495, "[ROCm] grid_sampler_2d backward: accumulate FP16/BF16 grid gradient in FP32 (changes values)", is the proposed fix. It is open and not merged (re-checked 2026-11-12). It touches only aten/src/ATen/native/cuda/GridSampler.cu, aten/src/ATen/native/cuda/GridSampler.cuh, docs/source/notes/hip.rst and test/test_nn.py; it is deliberately scoped to the ROCm build and leaves the CUDA build byte-identical. It does not touch intel/torch-xpu-ops, so it cannot cover XPU, and all three tests it adds are @onlyCUDA, so XPU CI will not surface this gap either.
• [CUDA][ROCm] grid_sampler_2d: keep 2D spatial offsets in the kernel's index type pytorch/pytorch#204493 and [ROCm] grid_sampler_2d backward: channel-parallel kernel, default on gfx1250 pytorch/pytorch#198032 are the other open members of the same grid-sampler series (32-bit offset overflow, channel-parallel ROCm kernel). Neither addresses the accumulator width and neither touches XPU.
• Found by an XPU alignment scan of the half-open window [2026-11-11T00:00:00Z, 2026-11-12T00:00:00Z) over pytorch/pytorch, frozen default-branch head 125472c0cbd6259715738ed1a80cd0240f914ad1.

Observed XPU behavior

float16   mask=[False, True] device=xpu:0 grad_grid=[1024.0, 0.0] exact=False
float16   mask=[True, True]  device=xpu:0 grad_grid=[1024.0, 0.0] exact=False
bfloat16  mask=[False, True] device=xpu:0 grad_grid=[128.0, 0.0] exact=False
bfloat16  mask=[True, True]  device=xpu:0 grad_grid=[128.0, 0.0] exact=False
float32   mask=[False, True] device=xpu:0 grad_grid=[2048.0, 0.0] exact=True
float32   mask=[True, True]  device=xpu:0 grad_grid=[2048.0, 0.0] exact=True

The float16 and bfloat16 rows are the exact values upstream reports for ROCm. Both output masks are affected, which places the loss in the grid-gradient accumulation itself rather than in the grad_input scatter (the [False, True] mask skips the scatter entirely).

Why [2048.0, 0.0] is the exact answer: with align_corners=False and inp_H = inp_W = 2, a grid coordinate of (0, 0) unnormalizes to ix = iy = 0.5, so all four bilinear weights are 0.25 and gix_mult = inp_W / 2 = 1. Per channel the x-gradient is -nw_val * (iy_se - iy) + ne_val * (iy_sw - iy) - sw_val * (iy - iy_ne) + se_val * (iy - iy_nw), which for the input column pattern [0, 1] is -0 * 0.5 + 1 * 0.5 - 0 * 0.5 + 1 * 0.5 = 1; over 2048 channels that is exactly 2048. The y-gradient cancels to 0. Because each channel adds two 0.5 taps, float16 stalls at 1024 (there the ULP is 1.0, so +0.5 rounds back to even) and bfloat16 stalls at 128 for the same reason.

Target path and root cause

aten/src/ATen/native/native_functions.yaml at the frozen head dispatches grid_sampler_2d_backward to XPU: grid_sampler_2d_backward_xpu, which this repository implements, so the call reaches XPU code and not a CPU fallback. Target-path evidence in the run: every result row reports device=xpu:0, torch.xpu.synchronize() is called before each value is read, and the float32 control on the same device and same kernel is exact, which rules out a harness or host/device artifact.

The 2-D bilinear backward kernel declares the grid-gradient accumulator in scalar_t and never widens it:

// src/ATen/native/xpu/sycl/GridSampler.cpp:365
  scalar_t gix_mult, giy_mult;

// src/ATen/native/xpu/sycl/GridSampler.cpp:390
    scalar_t gix = static_cast<scalar_t>(0), giy = static_cast<scalar_t>(0);

// src/ATen/native/xpu/sycl/GridSampler.cpp:448 (inside the per-channel loop)
        gix -= nw_val * (iy_se - iy) * gOut;

// src/ATen/native/xpu/sycl/GridSampler.cpp:473-474
    gGrid_ptr_NHW[0] = gix_mult * gix;
    gGrid_ptr_NHW[1] = giy_mult * giy;

The loop at src/ATen/native/xpu/sycl/GridSampler.cpp:395 walks all C channels and rounds the running sum back to scalar_t on every tap, so once the sum reaches 2^10 (float16) or 2^7 (bfloat16) relative to the tap size it stops advancing. The coordinate scale gix_mult is applied only after the loop, so widening it alone would not help. The forward kernel in the same file already computes in at::opmath_type<scalar_t> (src/ATen/native/xpu/sycl/GridSampler.cpp:65, 104), which is why only the backward grid gradient is wrong.

Two sibling kernels in the same file have the same narrow accumulator and are not separately reproduced here:

• 2-D bicubic backward: src/ATen/native/xpu/sycl/GridSampler.cpp:531-532, stored at 583-584.
• 3-D bilinear backward: src/ATen/native/xpu/sycl/GridSampler.cpp:1299-1300, stored at 1485-1487.

The 3-D bicubic backward kernel already accumulates in opmath_t and casts once at the store (src/ATen/native/xpu/sycl/GridSampler.cpp:1619-1621, 1665-1667), so the intended style is already present in this file.

Suggested fix: accumulate gix/giy (and giz) in at::opmath_type<scalar_t>, apply the coordinate multiplier in that type, and cast once when writing gGrid_ptr_*, matching lines 1619-1667. Upstream #204495 additionally keeps a per-channel partial for reduced precision to limit intermediate overflow and cancellation; that refinement is worth reviewing, but the accumulator width is the part that produces the numbers above.

Reproducer

import torch

for dtype in (torch.float16, torch.bfloat16, torch.float32):
    inp = (
        torch.tensor([[0.0, 1.0], [0.0, 1.0]])
        .expand(1, 2048, 2, 2)
        .to("xpu", dtype)
        .contiguous()
    )
    grad_output = torch.ones(1, 2048, 1, 1, device="xpu", dtype=dtype)
    grid = torch.zeros(1, 1, 1, 2, device="xpu", dtype=dtype)
    for mask in ([False, True], [True, True]):
        grad_grid = torch.ops.aten.grid_sampler_2d_backward(
            grad_output, inp, grid, 0, 0, False, mask
        )[1]
        torch.xpu.synchronize()
        print(dtype, mask, grad_grid.float().flatten().tolist(), "expected [2048.0, 0.0]")

This is the upstream snippet with device="cuda" changed to "xpu"; shapes, dtypes, the expand/contiguous layout, both output masks and the direct torch.ops.aten.grid_sampler_2d_backward call are unchanged, and the float32 row is an added control.

Environment

torch_version:   2.16.0.dev20261111+xpu
torch_path:      /opt/conda/lib/python3.13/site-packages/torch/__init__.py
python:          3.13.13 (/opt/conda/bin/python3.13)
xpu_available:   True
xpu_device_name: Intel(R) Arc(TM) Pro B60 Graphics

The reproducer ran unmodified as python3.13 -I -u with no network access and exited 0 in 3.5 s.

Ownership

The XPU kernel is independent code in this repository: grid_sampler_2d_backward_xpu (src/ATen/native/xpu/sycl/GridSampler.cpp:676) is dispatched directly from native_functions.yaml and shares nothing with aten/src/ATen/native/cuda/GridSampler.cuh. The only upstream fix in flight (#204495) is scoped to the ROCm build, leaves CUDA unchanged, and does not modify this repository, so no upstream change will fix XPU. Searches of pytorch/pytorch for XPU-scoped grid-sampler issues or PRs returned nothing.

No open intel/torch-xpu-ops issue covers this work. The closest existing items are different defects:

• [upstream_ut] XPU backward kernels missing input-validation error checks (replication_pad2d, grid_sample) #5508 (open) - XPU backward kernels missing input-validation error checks, including grid_sample; an error-checking gap, not accumulator width.
• grid_sample: implement mode='bicubic' with 5-D input on XPU #5542 and [release/2.15] NotImplementedError: grid_sampler(): bicubic interpolation with 5D input is not implemented for xpu #5452 (both closed) - grid_sample bicubic with 5-D input on XPU; a feature gap, unrelated to precision.

Related accumulator-width parity work already tracked separately: #6127 (segment_reduce fp16/bf16 accumulation) and #5722 (nll_loss2d fp16 global atomics). This issue is the same class of defect in a different kernel and does not overlap their code.

Suggested verification once fixed: the reproducer above must print [2048.0, 0.0] for float16 and bfloat16; upstream's test_grid_sample_2d_backward_low_precision_grid_grad (added by #204495, currently @onlyCUDA) is the natural regression test to mirror for XPU.