[Multimodal] Use GPU NVDEC for EPD encoder-only instance video media IO (#53675)

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

tests/config/test_multimodal_config.py

@@ -375,3 +375,85 @@ def test_vllm_config_runs_the_mm_processor_device_check():
         pytest.raises(ValueError, match="also runs the language model"),
     ):
         VllmConfig._validate_mm_processor_device(vllm_config)
+
+
+def _resolve_mm_video_decode_device(
+    *,
+    ec_role: ECRole | None,
+    mm_tensor_ipc: str = "torch_shm",
+    device: str | None = None,
+    video_kwargs: dict | None = None,
+    torchcodec_available: bool = True,
+) -> dict:
+    """Run processor-device then video-decode-device resolution and report
+    the resulting video media IO kwargs."""
+    mm_config = MultiModalConfig(
+        mm_processor_kwargs={} if device is None else {"device": device},
+        mm_tensor_ipc=mm_tensor_ipc,  # type: ignore[arg-type]
+        media_io_kwargs={} if video_kwargs is None else {"video": dict(video_kwargs)},
+    )
+    model_config = MagicMock(spec=ModelConfig)
+    model_config.multimodal_config = mm_config
+    vllm_config = MagicMock(spec=VllmConfig)
+    vllm_config.model_config = model_config
+    vllm_config.ec_transfer_config = (
+        None
+        if ec_role is None
+        else ECTransferConfig(ec_connector="ECExampleConnector", ec_role=ec_role)
+    )
+
+    with (
+        patch("vllm.platforms.current_platform.device_type", "cuda"),
+        patch(
+            "vllm.utils.import_utils.check_torchcodec_available",
+            side_effect=None if torchcodec_available else ImportError("torchcodec"),
+        ),
+    ):
+        VllmConfig._resolve_mm_processor_device(vllm_config)
+        VllmConfig._resolve_mm_video_decode_device(vllm_config)
+    return mm_config.media_io_kwargs.get("video", {})
+
+
+def test_auto_video_decode_uses_nvdec_on_encoder_instance():
+    """The processor runs on the accelerator there, so decoded frames should
+    stay on-device too."""
+    assert _resolve_mm_video_decode_device(ec_role="ec_producer") == {
+        "backend": "torchcodec",
+        "device": "cuda",
+    }
+
+
+@pytest.mark.parametrize("ec_role", [None, "ec_consumer", "ec_both"])
+def test_auto_video_decode_stays_on_cpu_off_encoder_instance(
+    ec_role: ECRole | None,
+):
+    """The processor stays on CPU off encode-only instances, so video
+    decoding should too."""
+    assert _resolve_mm_video_decode_device(ec_role=ec_role) == {}
+
+
+def test_auto_video_decode_follows_explicit_processor_device():
+    assert _resolve_mm_video_decode_device(ec_role="ec_producer", device="cpu") == {}
+    assert _resolve_mm_video_decode_device(ec_role="ec_producer", device="cuda") == {
+        "backend": "torchcodec",
+        "device": "cuda",
+    }
+
+
+def test_auto_video_decode_respects_explicit_media_io_kwargs():
+    assert _resolve_mm_video_decode_device(
+        ec_role="ec_producer", video_kwargs={"backend": "opencv"}
+    ) == {"backend": "opencv"}
+    assert _resolve_mm_video_decode_device(
+        ec_role="ec_producer", video_kwargs={"device": "cpu"}
+    ) == {"device": "cpu"}
+
+
+def test_auto_video_decode_skipped_without_torchcodec():
+    assert (
+        _resolve_mm_video_decode_device(
+            ec_role="ec_producer",
+            torchcodec_available=False,
+        )
+        == {}
+    )

tests/multimodal/media/test_video.py

@@ -532,6 +532,50 @@ class TestMergeKwargsGpuBackendPolicy:
     def test_unknown_backend_not_treated_as_gpu(self):
         assert not VIDEO_LOADER_REGISTRY.backend_requires_gpu("totally_unknown")
 
+    def test_strips_request_level_device(self):
+        """The decode device is a startup-only knob: a request must not move
+        decoding onto the GPU when the startup config did not opt in, nor off
+        it when it did."""
+        result = VideoMediaIO.merge_kwargs(
+            default_kwargs={"backend": "torchcodec"},
+            runtime_kwargs={"device": "cuda"},
+        )
+        assert "device" not in result
+
+        result = VideoMediaIO.merge_kwargs(
+            default_kwargs={"backend": "torchcodec", "device": "cuda"},
+            runtime_kwargs={"device": "cpu", "num_frames": 8},
+        )
+        assert result["device"] == "cuda"
+        assert result["num_frames"] == 8
+
+
+@pytest.mark.parametrize(
+    "default_kwargs",
+    [
+        {"backend": "torchcodec", "device": "cuda", "seek_mode": "approximate"},
+        {"backend": "pynvvideocodec", "hw_decoders": 2},
+    ],
+)
+def test_switching_backend_drops_stale_codec_options(default_kwargs):
+    """Codec-specific options from the static config must not leak into a
+    different codec backend selected per-request, where they would fail the
+    new backend's option validation."""
+    result = VideoMediaIO.merge_kwargs(
+        default_kwargs={**default_kwargs, "num_frames": 8},
+        runtime_kwargs={"backend": "opencv"},
+    )
+    assert result == {"backend": "opencv", "num_frames": 8}
+
+
+def test_same_backend_keeps_codec_options():
+    result = VideoMediaIO.merge_kwargs(
+        default_kwargs={"backend": "torchcodec", "device": "cuda"},
+        runtime_kwargs={"backend": "torchcodec", "num_frames": 8},
+    )
+    assert result["device"] == "cuda"
+    assert result["num_frames"] == 8
+
 
 @pytest.mark.parametrize("layout", ["nhwc", "nchw"])
 def test_pynvvc_frames_normalized_to_nhwc(layout: str):

tests/multimodal/test_hasher.py

@@ -137,6 +137,33 @@ def test_hash_collision_video_num_frames():
     )
 
 
+def test_hash_video_tensor_frames():
+    """Videos holding tensor frames (e.g. NVDEC-decoded) hash like
+    array-framed ones, from the original bytes without a D2H copy."""
+    source = b"x" * 100
+
+    def item_for_hash(frames):
+        metadata = {
+            "total_num_frames": 2,
+            "fps": 2.0,
+            "duration": 1.0,
+            "video_backend": "torchcodec",
+            "frames_indices": [0, 1],
+            "do_sample_frames": False,
+        }
+        video = MediaWithBytes((frames, metadata), source)
+        items = MultiModalDataParser()._parse_video_data([video])
+        return items.get_all_items_for_hash()[0]
+
+    np_frames = np.zeros((2, 8, 8, 3), dtype=np.uint8)
+    torch_frames = torch.zeros((2, 8, 8, 3), dtype=torch.uint8)
+
+    hasher = MultiModalHasher
+    assert hasher.hash_kwargs("blake3", video=item_for_hash(np_frames)) == (
+        hasher.hash_kwargs("blake3", video=item_for_hash(torch_frames))
+    )
+
+
 def test_hash_non_contiguous_array():
     arr = np.arange(24).reshape(4, 6).T
     assert not arr.flags.c_contiguous

tests/multimodal/test_parse.py

@@ -56,6 +56,28 @@ def test_frame_size_hwc_chw(frame):
     assert items.get_frame_size(0) == (W, H)
 
 
+def test_video_with_metadata_tensor_passthrough():
+    """Tensor frames pass through unchanged regardless of device: HF video
+    processors accept tensors, and device-resident frames (e.g. NVDEC-decoded)
+    must not be copied back to host."""
+    frames = torch.zeros((4, H, W, 3), dtype=torch.uint8)
+    video, metadata = MultiModalDataParser()._get_video_with_metadata(frames)
+
+    assert video is frames
+    assert metadata is None
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+def test_video_with_metadata_keeps_device_tensor():
+    """Device-resident frames (e.g. NVDEC-decoded) pass through as tensors,
+    so a device-side HF processor can consume them without a D2H copy."""
+    frames = torch.zeros((4, H, W, 3), dtype=torch.uint8, device="cuda")
+    video, metadata = MultiModalDataParser()._get_video_with_metadata(frames)
+
+    assert video is frames
+    assert metadata is None
+
+
 @pytest.mark.parametrize(
     "frames",
     [

tests/multimodal/test_video.py

@@ -11,6 +11,7 @@ from pathlib import Path
 import numpy as np
 import numpy.typing as npt
 import pytest
+import torch
 from transformers import AutoVideoProcessor
 from transformers.video_utils import VideoMetadata
 
@@ -173,9 +174,14 @@ def test_decode_video_imports_only_selected_backend(
     [
         (
             "torchcodec",
-            {"min_frames": 4, "num_ffmpeg_threads": 2, "seek_mode": "approximate"},
+            {
+                "min_frames": 4,
+                "num_ffmpeg_threads": 2,
+                "seek_mode": "approximate",
+                "device": "cuda",
+            },
             {"min_frames": 4},
-            {"num_ffmpeg_threads": 2, "seek_mode": "approximate"},
+            {"num_ffmpeg_threads": 2, "seek_mode": "approximate", "device": "cuda"},
         ),
         (
             "deepstream",
@@ -205,6 +211,11 @@ def test_video_backend_rejects_options_for_another_decoder():
     ):
         resolve_video_backend_kwargs("opencv", {"num_ffmpeg_threads": 2})
 
+    with pytest.raises(
+        ValueError, match="device is not supported by the 'opencv' backend"
+    ):
+        resolve_video_backend_kwargs("opencv", {"device": "cuda"})
+
 
 @pytest.mark.parametrize(
     ("backend", "error"),
@@ -1205,6 +1216,30 @@ def test_torchcodec_backend_rejects_frame_recovery(dummy_video_path):
         )
 
 
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
+def test_torchcodec_backend_cuda_decodes_on_gpu(dummy_video_path):
+    """With device="cuda", torchcodec decodes via NVDEC and keeps the frames
+    on the GPU instead of returning a host-side numpy array."""
+    pytest.importorskip("torchcodec")
+
+    with open(dummy_video_path, "rb") as f:
+        video_data = f.read()
+
+    loader = VIDEO_LOADER_REGISTRY.load("opencv")
+    frames, metadata = loader.load_bytes(
+        video_data, num_frames=8, backend="torchcodec", device="cuda"
+    )
+
+    assert isinstance(frames, torch.Tensor)
+    assert frames.device.type == "cuda"
+    assert frames.dtype == torch.uint8
+    assert frames.ndim == 4
+    assert frames.shape[3] == 3  # RGB
+    assert frames.shape[0] == 8
+    assert frames.shape[0] == len(metadata["frames_indices"])
+    assert metadata["video_backend"] == "torchcodec"
+
+
 def test_torchcodec_backend_returns_target_frames_not_keyframes():
     """Regression test: torchcodec must return the requested frames, not the
     GOP keyframe they seek back to.

vllm/config/vllm.py

@@ -1736,6 +1736,7 @@ class VllmConfig:
 
         self._resolve_mm_embedding_inputs()
         self._resolve_mm_processor_device()
+        self._resolve_mm_video_decode_device()
         self._validate_mm_processor_device()
 
         if self.use_v2_model_runner:
@@ -2623,6 +2624,61 @@ class VllmConfig:
             device_type,
         )
 
+    def _resolve_mm_video_decode_device(self) -> None:
+        """Default video decoding to torchcodec GPU backend for EPD encoder-only
+        instance if the mm processor runs on CUDA.
+
+        The processor consumes the decoded frames on-device in that case, so
+        keeping the frames on the GPU skips the host round-trip through the
+        CPU media path. An explicit codec/backend choice in
+        `--media-io-kwargs` is left alone, and the default is skipped where
+        torchcodec (or its FFmpeg runtime) is unavailable.
+        """
+        if self.model_config is None or self.model_config.multimodal_config is None:
+            return
+        mm_config = self.model_config.multimodal_config
+
+        ec_config = self.ec_transfer_config
+        # An EC producer that is not also a consumer runs no forward pass and
+        # allocates no KV cache, so frontend accelerator work has the device to
+        # itself.
+        if ec_config is None or not ec_config.is_encode_only:
+            return
+
+        from vllm.platforms import current_platform
+
+        device_type = current_platform.device_type
+        if (
+            device_type != "cuda"
+            or mm_config.get_mm_processor_device_type() != device_type
+        ):
+            return
+
+        # User set video backend or device explicitly
+        video_kwargs = mm_config.media_io_kwargs.setdefault("video", {})
+        if "backend" in video_kwargs or "device" in video_kwargs:
+            return
+
+        from vllm.utils.import_utils import check_torchcodec_available
+
+        try:
+            check_torchcodec_available()
+        except (ImportError, RuntimeError):
+            # torchcodec is not installed, or is installed without a usable
+            # FFmpeg runtime (it raises rather than returning False).
+            logger.info_once(
+                "EPD encoder instance: keeping CPU video decoding because "
+                "torchcodec is not available (needs a CUDA build with FFmpeg)."
+            )
+            return
+
+        video_kwargs["backend"] = "torchcodec"
+        video_kwargs["device"] = device_type
+        logger.info_once(
+            "EPD encoder instance: decoding video with NVDEC (torchcodec device=%s).",
+            device_type,
+        )
+
     def _validate_mm_processor_device(self) -> None:
         """Hand the EC config to `MultiModalConfig`, which owns the rule."""
         model_config = self.model_config

vllm/multimodal/hasher.py

@@ -89,8 +89,11 @@ class MultiModalHasher:
                 )
             return cls.iter_item_to_bytes("image", obj.original_bytes)
 
-        if isinstance(obj, MediaWithBytes) and isinstance(obj.media, np.ndarray):
+        if isinstance(obj, MediaWithBytes) and isinstance(
+            obj.media, (np.ndarray, torch.Tensor)
+        ):
             frames = obj.media
+            # Both np.ndarray and torch.Tensor expose .nbytes.
             if frames.nbytes < len(obj.original_bytes):
                 return cls.iter_item_to_bytes("video", frames)
             return cls.iter_item_to_bytes("video", obj.original_bytes)

vllm/multimodal/media/video.py

@@ -19,17 +19,23 @@ from vllm.utils.sparse_utils import (
     safe_to_dense,
 )
 
-from ..video import VIDEO_LOADER_REGISTRY
+from ..video import VIDEO_LOADER_REGISTRY, DecodedFrames
+from ..video_decoders import incompatible_backend_options
 from .base import MediaIO, MediaWithBytes
 from .image import MAGIC_NUMPY_PREFIX, ImageMediaIO
 
 logger = init_logger(__name__)
 
 
-class VideoMediaIO(MediaIO[MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]]):
+class VideoMediaIO(MediaIO[MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]]):
     """Configuration values can be user-provided either by --media-io-kwargs or
     by the runtime API field "media_io_kwargs". Ensure proper validation and
     error handling.
+
+    Use --media-io-kwargs '{"video": {"backend": "torchcodec",
+    "device": "cuda"}}' to decode with NVDEC and keep frames on the GPU,
+    which pairs with a device-side multi-modal processor
+    (--mm-processor-device, EPD encode-only instances).
     """
 
     @classmethod
@@ -43,6 +49,9 @@ class VideoMediaIO(MediaIO[MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]]):
             runtime_kwargs = dict(runtime_kwargs)
             runtime_kwargs.pop("hw_decoders", None)
             runtime_kwargs.pop("pool_size", None)
+            # The decode device determines which device the decoded frames
+            # (and thus the processor pipeline) live on; startup-only.
+            runtime_kwargs.pop("device", None)
 
             # Block request-level selection of GPU video backends that
             # were not configured (and VRAM-reserved) at startup.
@@ -70,6 +79,14 @@ class VideoMediaIO(MediaIO[MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]]):
                 merged.pop("fps", None)
             elif "fps" in runtime_kwargs and "num_frames" not in runtime_kwargs:
                 merged.pop("num_frames", None)
+
+            requested_backend = runtime_kwargs.get("backend")
+            static_backend = (default_kwargs or {}).get("backend")
+            if requested_backend and requested_backend != static_backend:
+                for name in incompatible_backend_options(
+                    static_backend, requested_backend
+                ):
+                    merged.pop(name, None)
         return merged
 
     def __init__(
@@ -101,7 +118,7 @@ class VideoMediaIO(MediaIO[MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]]):
 
     def load_bytes(
         self, data: bytes
-    ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]:
+    ) -> MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]:
         video = self.video_loader.load_bytes(
             data, num_frames=self.num_frames, **self.kwargs
         )
@@ -109,7 +126,7 @@ class VideoMediaIO(MediaIO[MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]]):
 
     def load_base64(
         self, media_type: str, data: str
-    ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]:
+    ) -> MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]:
         if media_type.lower() == "video/jpeg":
             load_frame = partial(
                 self.image_io.load_base64,
@@ -177,7 +194,7 @@ class VideoMediaIO(MediaIO[MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]]):
 
     def load_file(
         self, filepath: Path
-    ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]:
+    ) -> MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]:
         with filepath.open("rb") as f:
             data = f.read()
 

vllm/multimodal/parse.py

@@ -36,6 +36,7 @@ from .inputs import (
     VideoItem,
 )
 from .media import MediaWithBytes
+from .video import DecodedFrames
 
 _T = TypeVar("_T")
 _I = TypeVar("_I")
@@ -653,7 +654,7 @@ class MultiModalDataParser:
     def _get_video_with_metadata(
         self,
         video: VideoItem,
-    ) -> tuple[np.ndarray | MediaWithBytes[np.ndarray], dict[str, Any] | None]:
+    ) -> tuple[DecodedFrames | MediaWithBytes[DecodedFrames], dict[str, Any] | None]:
         if isinstance(video, MediaWithBytes):
             new_video, metadata = self._get_video_with_metadata(video.media)
             return MediaWithBytes(new_video, video.original_bytes), metadata
@@ -661,10 +662,11 @@ class MultiModalDataParser:
             return video
         if isinstance(video, list):
             return np.array(video), None
-        if isinstance(video, np.ndarray):
+        if isinstance(video, (np.ndarray, torch.Tensor)):
+            # Tensors pass through untouched: HF video processors accept them
+            # directly, and device tensors (e.g. NVDEC-decoded frames) must
+            # stay on-device to avoid a D2H round-trip.
             return video, None
-        if isinstance(video, torch.Tensor):
-            return video.numpy(), None
 
         assert_never(video)
 
@@ -769,9 +771,9 @@ class MultiModalDataParser:
             data_items = data  # type: ignore[assignment]
 
         new_videos = list[
-            np.ndarray
-            | MediaWithBytes[np.ndarray]
-            | tuple[np.ndarray | MediaWithBytes[np.ndarray], dict[str, Any]]
+            DecodedFrames
+            | MediaWithBytes[DecodedFrames]
+            | tuple[DecodedFrames | MediaWithBytes[DecodedFrames], dict[str, Any]]
             | None
         ]()
         metadata_lst: list[dict[str, Any] | None] = []

vllm/multimodal/video.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import math
 from abc import abstractmethod
-from typing import Any, ClassVar, cast
+from typing import Any, ClassVar, TypeAlias, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -28,6 +28,10 @@ except ImportError:
 
 logger = init_logger(__name__)
 
+DecodedFrames: TypeAlias = npt.NDArray | torch.Tensor
+"""Decoded video frames: a host ``np.ndarray``, or a device ``torch.Tensor``
+when a GPU decoding codec is used (e.g. torchcodec with ``device="cuda"``)."""
+
 
 class VideoLoaderRegistry(ExtensionManager):
     def __init__(self) -> None:
@@ -164,8 +168,13 @@ class VideoLoader:
         cls,
         data: bytes,
         **kwargs,
-    ) -> tuple[npt.NDArray, dict[str, Any]]:
-        """Load video frames from bytes and return (frames_array, metadata_dict)."""
+    ) -> tuple[DecodedFrames, dict[str, Any]]:
+        """Load video frames from bytes and return (frames, metadata_dict).
+
+        ``frames`` is a CPU ``np.ndarray`` unless a GPU decoding codec is
+        used (e.g. torchcodec with ``device="cuda"``), in which case it is
+        a ``torch.Tensor`` living on that device.
+        """
         raise NotImplementedError
 
     @classmethod
@@ -239,7 +248,7 @@ class VideoBackend(VideoLoader):
         *,
         backend: VideoDecoderBackend = "opencv",
         **kwargs: Any,
-    ) -> tuple[npt.NDArray, dict[str, Any]]:
+    ) -> tuple[DecodedFrames, dict[str, Any]]:
         """Load sampled frames from raw video bytes.
 
         Args:
@@ -265,6 +274,11 @@ class VideoBackend(VideoLoader):
                   creation at the cost of relying on the file's metadata. See
                   https://meta-pytorch.org/torchcodec/stable/generated_examples/decoding/approximate_mode.html
                   for details.
+                - ``device`` (TorchCodec): ``"cpu"`` (default) decodes on the
+                  host and returns a ``np.ndarray``; ``"cuda"`` decodes with
+                  NVDEC and returns the frames as a CUDA ``torch.Tensor``,
+                  which a device-side HF video processor can consume without
+                  a host round-trip.
                 - ``hw_decoders`` (PyNvVideoCodec): maximum number of
                   concurrent decoder slots. Defaults to 2 and must be a
                   positive integer.
@@ -272,7 +286,8 @@ class VideoBackend(VideoLoader):
                   size and pool acquisition timeout in seconds.
 
         Returns:
-            Tuple of ``(frames_array, metadata_dict)``.
+            Tuple of ``(frames, metadata_dict)``, where ``frames`` is a
+            CPU ``np.ndarray`` unless TorchCodec decodes on ``device="cuda"``.
         """
         target = VideoTargetMetadata(
             num_frames=num_frames, fps=fps, max_duration=max_duration

vllm/multimodal/video_decoders/__init__.py

@@ -19,6 +19,7 @@ _BACKEND_OPTION_DEFAULTS: dict[str, dict[str, Any]] = {
     "torchcodec": {
         "num_ffmpeg_threads": 0,
         "seek_mode": "exact",
+        "device": "cpu",
     },
     PYNVVIDEOCODEC_VIDEO_BACKEND: {
         "hw_decoders": PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,
@@ -40,6 +41,17 @@ def _get_backend_option_defaults(backend: str) -> dict[str, Any]:
         ) from None
 
 
+# TODO(Isotr0py): Make kwargs structured.
+def incompatible_backend_options(
+    old_backend: str | None,
+    new_backend: str,
+) -> set[str]:
+    """Option names accepted by ``old_backend`` but not by ``new_backend``."""
+    old = _BACKEND_OPTION_DEFAULTS.get(old_backend or "", {})
+    new = _BACKEND_OPTION_DEFAULTS.get(new_backend, {})
+    return old.keys() - new.keys()
+
+
 def resolve_video_backend_kwargs(
     backend: str,
     kwargs: dict[str, Any],
@@ -103,5 +115,6 @@ __all__ = [
     "VideoTargetMetadata",
     "check_frame_pixel_limit",
     "decode_video",
+    "incompatible_backend_options",
     "resolve_video_backend_kwargs",
 ]

vllm/multimodal/video_decoders/torchcodec.py

@@ -5,7 +5,10 @@ from typing import Literal
 
 import numpy as np
 import numpy.typing as npt
+import torch
 
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
 from vllm.utils.import_utils import PlaceholderModule, check_torchcodec_available
 
 from .base import (
@@ -21,6 +24,8 @@ except (ImportError, RuntimeError):
         "decoders.VideoDecoder"
     )
 
+logger = init_logger(__name__)
+
 
 def decode_torchcodec(
     loader_cls,
@@ -30,12 +35,14 @@ def decode_torchcodec(
     *,
     num_ffmpeg_threads: int = 0,
     seek_mode: Literal["exact", "approximate"] = "exact",
-) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:
+    device: str = "cpu",
+) -> tuple[npt.NDArray | torch.Tensor, VideoSourceMetadata, list[int], list[int]]:
     check_torchcodec_available()
     decoder = TorchCodecVideoBackendMixin.make_torchcodec_decoder(
         data,
         num_ffmpeg_threads=num_ffmpeg_threads,
         seek_mode=seek_mode,
+        device=device,
     )
     check_frame_pixel_limit(
         decoder.metadata.width or 0,
@@ -48,7 +55,7 @@ def decode_torchcodec(
         source=source, target=target, **sampling_kwargs
     )
     frames, valid = TorchCodecVideoBackendMixin.decode_torchcodec_frames(
-        decoder, frame_idx
+        decoder, frame_idx, device=device
     )
     return frames, source, frame_idx, valid
 
@@ -67,7 +74,19 @@ class TorchCodecVideoBackendMixin:
         *,
         num_ffmpeg_threads: int = 0,
         seek_mode: Literal["exact", "approximate"] = "exact",
+        device: str = "cpu",
     ) -> "VideoDecoder":
+        torch_device = torch.device(device)
+        if torch_device.type == "cuda" and not current_platform.is_cuda():
+            raise ValueError(
+                f"torchcodec video decoding on device {device!r} requires "
+                "a CUDA-capable platform."
+            )
+        elif torch_device.type not in ("cpu", "cuda"):
+            raise ValueError(
+                f"torchcodec video decoding only supports 'cpu' and 'cuda' "
+                f"devices, got {device!r}."
+            )
         # NHWC matches the (num_frames, H, W, 3) uint8 RGB layout the rest
         # of the pipeline expects, avoiding a transpose.
         return VideoDecoder(
@@ -75,6 +94,7 @@ class TorchCodecVideoBackendMixin:
             dimension_order="NHWC",
             num_ffmpeg_threads=num_ffmpeg_threads,
             seek_mode=seek_mode,
+            device=device,
         )
 
     @staticmethod
@@ -91,10 +111,23 @@ class TorchCodecVideoBackendMixin:
     def decode_torchcodec_frames(
         decoder: "VideoDecoder",
         frame_indices: list[int],
-    ) -> tuple[npt.NDArray, list[int]]:
+        *,
+        device: str = "cpu",
+    ) -> tuple[npt.NDArray | torch.Tensor, list[int]]:
         """Decode the requested indices in one batched, index-exact call."""
         if not frame_indices:
             return np.empty((0,), dtype=np.uint8), []
         # Note: torchcodec releases the GIL for the entire call
         batch = decoder.get_frames_at(frame_indices)
-        return batch.data.numpy(), list(frame_indices)
+        frames = batch.data
+        if (
+            torch.device(device).type != "cpu"
+            and (status := getattr(decoder, "cpu_fallback", None)) is not None
+            and bool(status)
+        ):
+            logger.warning_once(
+                "torchcodec could not use NVDEC for this video and "
+                "decoded on CPU instead; check codec support and "
+                "libnvcuvid."
+            )
+        return frames, list(frame_indices)