Fix CUDA decode_jpeg writing into memory the caller's stream is still using

[karkuspeter]

The bug

CUDAJpegDecoder::decode_images() gets its output tensors from prepare_buffers(), which calls
torch::stable::empty() while the caller's stream is current. It then submits the nvJPEG work on
the decoder's own stream, taken from torch's stream pool in the constructor. Nothing makes that
private stream wait on the caller's stream first.

That breaks the caching allocator's contract. The allocator hands a freed block straight back out for
reuse on the stream it was freed on, because every access to a block allocated on a stream is expected
to be ordered on that stream. So when the caller frees a tensor while kernels reading it are still
queued — an ordinary temporary going out of scope — the allocator can hand that same block to
prepare_buffers() on the next decode_jpeg() call, and nvJPEG writes the decoded image into it from
an unordered stream, possibly long before those queued kernels run.

The damage does not show up in the decoded image. It shows up in whatever unrelated tensor was
overwritten: silently wrong values, and NaNs once the clobbered bytes are read as floats. No threads
and no explicit streams are needed in user code, only a caller with GPU work in flight, which is the
normal case when decoding is interleaved with compute.

decode_jpegs_cuda() already closes the other half of this loop: after decode_images() returns, it
makes the caller's current stream wait on the decoder's stream, so the caller cannot read the image
before it has been written. The missing half is the decoder waiting for the caller before it writes.

This predates the stable-ABI port (#9533). v0.28.0 and earlier allocate the outputs with
torch::empty() and decode on at::cuda::getStreamFromPool() in exactly the same way, so it is not a
regression, and released versions are affected.

The fix

Record an event on the caller's current stream once the buffers exist, and have the decoder's stream
wait on it, through the syncStreams() helper already in this file. It is the mirror image of the sync
that already runs after the decode.

Reproduction

Single-threaded, no explicit streams, public APIs only. Each iteration fills a few tensors with known
values and synchronizes, queues a long-running kernel followed by one reduction per tensor, frees the
tensors while those reductions are still queued, calls decode_jpeg, then synchronizes and checks the
reductions against their exact expected value. Nothing in that sequence is racy: a wrong sum can only
mean the decoder wrote into memory whose readers had not run yet.

repro.py

"""decode_jpeg(device="cuda") writes into memory the caller's stream is still using.

CUDAJpegDecoder allocates its output tensors with torch::empty while the *caller's*
stream is current, then submits the nvJPEG work on a private stream taken from
torch's stream pool without first making that stream wait on the caller's stream.
The caching allocator recycles a block as soon as it is freed on the stream that
owns it, because all later use is supposed to be ordered on that same stream. So
the decoder is regularly handed a block that kernels queued on the caller's stream
are still going to read, and nvJPEG overwrites it before those kernels run.

Single threaded and deterministic. Each iteration:

  1. fills N tensors with known values and synchronizes, so the data is resident;
  2. queues a long-running kernel, then queues one reduction per tensor;
  3. frees the tensors, while those reductions are still queued;
  4. calls decode_jpeg, which is handed one of the freed blocks;
  5. synchronizes and checks the reductions against their exact expected value.

Nothing here is racy by itself: step 5 only sees a wrong number if the decoder
wrote into memory that step 2's kernels had not finished reading.

    python repro.py                        # corrupted > 0
    python repro.py --decode none          # control: no decode          -> 0
    python repro.py --decode cpu           # control: CPU decode         -> 0
    python repro.py --stall 0              # control: nothing in flight  -> 0
    python repro.py --caller-stream side   # workaround: side stream     -> 0
    python repro.py --decode mimic         # same stream pattern, no nvJPEG
"""

from __future__ import annotations

import argparse
import contextlib

import torch
import torchvision
from torchvision.io import ImageReadMode, decode_jpeg, encode_jpeg

HEIGHT, WIDTH = 1080, 1920
# float32 tensors with exactly the decode output's byte count (3 * H * W uint8), so
# the allocator can hand the decoder a block that was just freed.
NUMEL = 3 * HEIGHT * WIDTH // 4


def make_jpeg(path: str | None) -> torch.Tensor:
    """A 1920x1080 JPEG, either read from `path` or synthesized here."""
    if path is not None:
        return torch.frombuffer(bytearray(open(path, "rb").read()), dtype=torch.uint8)
    rows = torch.linspace(0, 1, HEIGHT).unsqueeze(1)
    cols = torch.linspace(0, 1, WIDTH).unsqueeze(0)
    plane = ((torch.sin(12 * cols) * torch.cos(9 * rows) + 1) * 127).to(torch.uint8)
    return encode_jpeg(plane.expand(3, HEIGHT, WIDTH).contiguous(), quality=90)


def decode(jpeg: torch.Tensor, device: torch.device, kind: str, side: torch.cuda.Stream) -> None:
    if kind == "mimic":
        # The decoder's stream pattern with nvJPEG removed: allocate while the
        # caller's stream is current, write from a stream out of the same pool,
        # host-sync that stream. No torchvision code involved.
        buffer = torch.empty((3, HEIGHT, WIDTH), dtype=torch.uint8, device=device)
        with torch.cuda.stream(side):
            buffer.fill_(0x7F)
        side.synchronize()
    else:
        decode_jpeg(jpeg, mode=ImageReadMode.RGB, device=device if kind == "cuda" else "cpu")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--decode", default="cuda", choices=("cuda", "cpu", "none", "mimic"))
    parser.add_argument("--iterations", type=int, default=20)
    parser.add_argument("--victims", type=int, default=4, help="tensors read by queued kernels")
    parser.add_argument(
        "--stall",
        type=int,
        default=200_000_000,
        help="cycles of GPU work queued ahead of the reductions (~100ms on H100); 0 disables",
    )
    parser.add_argument(
        "--caller-stream",
        default="default",
        choices=("default", "side"),
        help="stream that is current when decode_jpeg is called",
    )
    parser.add_argument("--jpeg", help="1920x1080 JPEG file to decode; synthesized if omitted")
    args = parser.parse_args()

    device = torch.device("cuda", torch.cuda.current_device())
    print(
        f"torch {torch.__version__}  torchvision {torchvision.__version__}  "
        f"{torch.cuda.get_device_name(device)}  allocator {torch.cuda.get_allocator_backend()}\n"
        f"decode={args.decode} iterations={args.iterations} victims={args.victims} "
        f"stall={args.stall} caller_stream={args.caller_stream}",
        flush=True,
    )

    jpeg = make_jpeg(args.jpeg)
    side = torch.cuda.Stream(device=device)
    caller_scope = (
        torch.cuda.stream(torch.cuda.Stream(device=device))
        if args.caller_stream == "side"
        else contextlib.nullcontext()
    )
    if args.decode != "none":
        # One-time decoder setup, kept out of the measurement.
        decode(jpeg, device, args.decode, side)
        torch.cuda.synchronize()

    expected = float(sum((index + 1) * NUMEL for index in range(args.victims)))
    corrupted = 0
    first = ""
    for _ in range(args.iterations):
        victims = [
            torch.full((NUMEL,), float(index + 1), device=device)
            for index in range(args.victims)
        ]
        torch.cuda.synchronize()  # victim data is resident on the device

        if args.stall:
            torch.cuda._sleep(args.stall)
        total = torch.zeros((), dtype=torch.float64, device=device)
        for victim in victims:
            total = total + victim.sum(dtype=torch.float64)  # queued read
        victims.clear()  # freed while the reductions are still queued

        if args.decode != "none":
            with caller_scope:
                decode(jpeg, device, args.decode, side)

        torch.cuda.synchronize()
        got = float(total)
        if got != expected:
            corrupted += 1
            if not first:
                first = f"first mismatch: got {got:.6g}, expected {expected:.6g}"

    print(f"corrupted={corrupted}/{args.iterations}", flush=True)
    if first:
        print(first, flush=True)


if __name__ == "__main__":
    main()

On main:

$ python repro.py --iterations 200
torch 2.12.0.dev20260408+cu128  torchvision 0.29.0a0+fc872e4  NVIDIA H100 80GB HBM3  allocator native
decode=cuda iterations=200 victims=4 stall=200000000 caller_stream=default
corrupted=199/200
first mismatch: got nan, expected 1.5552e+07

The corrupted sums are arbitrary: NaN, 5.3e+44, small integers, whatever the decoded pixels happen
to be when read as floats.

Verification

Built from source on an H100 80GB, torch 2.12.0.dev20260408+cu128, main at fc872e4, native
allocator. Each leg below was run on a build of this exact tree, and on a build with only the C++
change reverted.

The reproducer and its controls, 200 iterations each:

	main	this PR
repro.py	199/200 corrupted	0/200
repro.py --decode none	0/200	0/200
repro.py --decode cpu	0/200	0/200
repro.py --stall 0	0/200	0/200
repro.py --caller-stream side	0/200	0/200

The controls matter as much as the failure: with no decode, with a CPU decode, or with nothing in
flight on the caller's stream, the harness is clean, so the corruption is specifically the CUDA decode
writing over the caller's queued reads. --caller-stream side calls decode_jpeg with a different
stream current from the one the freed tensors belonged to, so the outputs come from that stream's pool,
nothing aliases and nothing breaks — which is what makes the allocation-time stream the culprit.

The script's --decode mimic leg replaces the decode with the same allocate-here/write-there pattern
written by hand in Python — allocate while the caller's stream is current, write from another stream,
host-sync that stream, no torchvision involved. It corrupts at the same rate on either build (19/20 in
a 20-iteration run), which is what pins the mechanism to the stream pattern rather than to nvJPEG.

The regression test added here, five consecutive runs each way, no flakiness in either direction:

main:     1 failed, 688 deselected   (x5)
          FAILED test/test_image.py::test_decode_jpegs_cuda_stream_ordering
          AssertionError: assert 833739888.0 == 62208000.0     (the wrong value varies)
this PR:  1 passed, 688 deselected   (x5)

No regressions in the existing tests. pytest test/test_image.py -k "jpeg or cuda" selects 93 tests
including the new one:

main:     1 failed, 92 passed, 8 xfailed     (only the new test fails)
this PR:  93 passed, 8 xfailed

No measurable cost. Decode throughput for a 1920x1080 JPEG, A/B/A over three builds, 5x100 decodes
each:

patched    1.154 ms/image
unpatched  1.146 ms/image
patched    1.151 ms/image

black and clang-format are clean on the changed files; flake8 reports only two pre-existing E231s
elsewhere in test/test_image.py.

Notes

• An alternative fix would allocate the outputs while the decoder's stream is current and call
  record_stream(caller_stream) on them before returning, so the calling thread would never wait for
  its own queued work. record_stream is not exposed through the stable ABI this file now uses
  (torch/csrc/stable/c/shim.h offers get/set current stream, stream-from-pool and stream
  synchronize), so the event wait is the smaller change. Happy to go the other way if you prefer.
• This does not change the op's synchronization semantics: decode_images() already ends with
  cudaStreamSynchronize(stream), so the decode is host-synchronous either way. The patch only adds
  the missing ordering edge.
• The test uses torch.cuda._sleep to keep the caller's stream busy, and skips if the allocator
  happened not to hand the decoder one of the freed blocks, so it cannot fail spuriously. Happy to
  replace the private helper with something else.
• Out of scope here, but worth flagging: the encode path has the same class of problem. encode_jpeg generates corrupted JPEG data when using CUDA #9060 reports
  corrupted output from encode_jpeg on CUDA, and while Encode jpeg cuda sync #8929 added a sync there, the encoder captures
  current_stream in its constructor, so a later call can end up synchronizing against a stream that
  is no longer the caller's.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/9606

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[meta-cla]

Hi @karkuspeter!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[NicolasHug]

Thanks for the PR @karkuspeter , I'm not going to land it because I'm going to deprecate the torchvision decoders in the next version (~next month). The new jpeg cuda decoder is in TorchCodec now, which I just released in TorchCodec 0.16 https://meta-pytorch.org/torchcodec/stable/generated/torchcodec.decoders.decode_jpeg.html#torchcodec.decoders.decode_jpeg

Let me know if that one still has the same issue - happy to consider a fix in TorchCodec instead.

[karkuspeter]

Thanks — I tested TorchCodec 0.16. It still corrupts, but not because of the wrapper: that design is
right, and what's left is an nvJPEG bug no wrapper can avoid without a host barrier.

One harness for both, on an H100 (driver 575.57.08, torch 2.13.0+cu129, released wheels). Each
iteration queues ~100 ms of work and then one reduction per tensor on the caller's stream, frees the
tensors while those reductions are still queued, and decodes. The decoder is handed one of the freed
blocks in 135 of 200 iterations; corrupted counts iterations where a reduction returned the wrong
number.

decoder	JPEG	nvJPEG path	corrupted	decode call
torchvision 0.28	baseline	hardware	134/200	1.2 ms
torchvision 0.28	progressive	software	134/200	3.3 ms
torchcodec 0.16	baseline	hardware	134/200	100.2 ms
torchcodec 0.16	progressive	software	0/200	100.7 ms

Draining the caller's stream before the decode makes every row 0/200.

The progressive rows are the discriminator: that JPEG takes nvJPEG's software path, which orders
correctly, so TorchCodec is clean there while torchvision still corrupts. This PR therefore fixes a
real wrapper defect that TorchCodec does not have. The call times say the same thing — torchvision
returns in 1.2 ms with ~100 ms queued on the caller's stream, because it decodes on a private stream
nothing orders against it, while TorchCodec waits, because it submits on the caller's stream.

What remains on baseline JPEGs is nvJPEG itself. In a standalone CUDA program — no framework, one
cudaMalloc'd destination reused for the whole run, a spin kernel and a checksum kernel queued ahead
of the decode on the same stream — nvjpegDecodeBatched on a hardware-backend handle writes the
destination before that queued read runs, 100 times out of 100. With cudaStreamSynchronize first:
0/100. With a cudaStreamWaitEvent on the submission stream instead: still 100/100. With that same
event waited for on the host: 0/100. So the write happens when the call is made, not when the
submission stream comes free, and a host barrier is the only defence available. nvjpegDecode and the
software backend are both clean. Reproduced on nvJPEG 12.3.5, 12.4.0.76, 13.0.1 and 13.2.1; an A100
was clean over 300 iterations, so CI on other hardware may not see it. Filing it with NVIDIA
separately; DALI carries the same workaround for nvImageCodec in NVIDIA/DALI#5408.

One correction to this PR. It is verified (199/200 corrupted before, 0/200 after, plus the regression
test), but the rationale I gave is incomplete: the event I record is only effective because the
unconditional cudaStreamSynchronize on the decoder's private stream a few lines below turns it into
a host barrier. Per the numbers above, an event alone would not have been enough. Worth knowing if that
sync is ever removed to make decode asynchronous.

For the TorchCodec side I've opened meta-pytorch/torchcodec#1634, which host-synchronizes the caller's
stream in decode_batched_hardware before the first nvjpegDecodeBatched. Hardware path only, and
free in wall-clock terms because decode_images already host-synchronizes that stream before
returning — median 1080p decode 1.186 ms with the barrier and 1.183 ms without, and the regression test
there fails 3 runs out of 3 without it and passes 5 out of 5 with it. Both reproducers are inline in
that PR; the CUDA one needs no input data, it encodes its own JPEG.

This PR is your call, and either way is fine by me. 0.28 silently corrupts unrelated GPU memory today
on both nvJPEG paths, and it surfaces as NaNs somewhere unrelated rather than as anything pointing at
the decoder, so landing it would cover the releases before the decoders go away. If you'd rather not
touch code you're deprecating, closing it is fine, and a line in the deprecation note saying CUDA
decode in ≤0.28 can corrupt concurrent work would help anyone who can't move yet.

[NicolasHug]

thanks, I'll close this and follow-up on meta-pytorch/torchcodec#1634