fix(vlm): preserve Kimi-K3 GPU JPEG accuracy (#34163)

python/sglang/srt/disaggregation/encode_server.py

@@ -647,7 +647,12 @@ class MMEncoder:
             return data
         try:
             if modality == Modality.IMAGE:
-                img, _ = load_image(data, False)
+                gpu_image_decode = (
+                    "nvjpeg_fancy"
+                    if self.use_image_processor_gpu and self.model_type == "kimi_k3"
+                    else False
+                )
+                img, _ = load_image(data, gpu_image_decode)
                 if (
                     discard_alpha_channel
                     and not isinstance(img, torch.Tensor)

python/sglang/srt/multimodal/processors/kimi_k3.py

@@ -282,7 +282,10 @@ class KimiK3GPUProcessorWrapper(KimiGPUProcessorWrapper):
 
 class KimiK3ImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor):
     models = [KimiK3ForConditionalGeneration]
-    gpu_image_decode = True
+    # K3 accuracy is sensitive to the chroma upsampling used for common 4:2:0
+    # JPEG inputs. This mode uses interpolated nvJPEG upsampling when the K3
+    # image dependency is installed and otherwise falls back to PIL.
+    gpu_image_decode = "nvjpeg_fancy"
     prefer_tokenized_input = True
     precompute_hash_before_cpu_transfer = True
     auto_mm_processor_worker_num = 2

python/sglang/srt/utils/common.py

@@ -1643,9 +1643,12 @@ class VideoData:
 
 
 image_extension_names = (".png", ".jpg", ".jpeg", ".webp", ".gif")
+GPUImageDecodeMode = Union[bool, Literal["nvjpeg_fancy"]]
 
 
-def is_jpeg_with_cuda(image_bytes: bytes = b"", gpu_image_decode: bool = True) -> bool:
+def is_jpeg_with_cuda(
+    image_bytes: bytes = b"", gpu_image_decode: GPUImageDecodeMode = True
+) -> bool:
     """
     Check three conditions:
     1. whether CUDA is available.
@@ -1659,10 +1662,19 @@ def is_jpeg_with_cuda(image_bytes: bytes = b"", gpu_image_decode: bool = True) -
     return False
 
 
+@lru_cache(maxsize=16)
+def _warn_fancy_jpeg_fallback(error: str) -> None:
+    logger.warning(
+        "High-fidelity GPU JPEG decode is unavailable; falling back to PIL. "
+        "Install the Kimi-K3 serving image or NVIDIA nvImageCodec. Error: %s",
+        error,
+    )
+
+
 def _load_image(
     image_bytes: bytes = b"",
     image_file: str = "",
-    gpu_image_decode: bool = True,
+    gpu_image_decode: GPUImageDecodeMode = True,
 ) -> Union[torch.Tensor, Image.Image]:
     """
     Try to decode JPEG with nvJPEG on GPU and return a torch device tensor,
@@ -1673,19 +1685,29 @@ def _load_image(
         image_bytes = get_image_bytes(image_file)
     if is_jpeg_with_cuda(image_bytes, gpu_image_decode):
         try:
+            if gpu_image_decode == "nvjpeg_fancy":
+                from sglang.srt.utils.nvjpeg_decoder import (
+                    decode_jpeg_with_fancy_upsampling,
+                )
+
+                return decode_jpeg_with_fancy_upsampling(image_bytes)
             encoded_image = torch.frombuffer(image_bytes, dtype=torch.uint8)
             image_tensor = decode_jpeg(encoded_image, device="cuda")
             return image_tensor
         except Exception as e:
-            logger.warning(
-                f"Failed to decode JPEG on GPU, falling back to CPU. Error: {e}"
-            )
+            if gpu_image_decode == "nvjpeg_fancy":
+                _warn_fancy_jpeg_fallback(f"{type(e).__name__}: {e}")
+            else:
+                logger.warning(
+                    "Failed to decode JPEG on GPU, falling back to CPU. Error: %s",
+                    e,
+                )
     return Image.open(BytesIO(image_bytes))
 
 
 def load_image(
     image_file: Union[Image.Image, str, ImageData, bytes],
-    gpu_image_decode: bool = True,
+    gpu_image_decode: GPUImageDecodeMode = True,
 ) -> tuple[Union[torch.Tensor, Image.Image], Optional[tuple[int, int]]]:
     """
     Load image from multiple input formats, including:

python/sglang/srt/utils/nvjpeg_decoder.py

@@ -0,0 +1,86 @@
+"""High-fidelity nvJPEG decoding for multimodal image processors."""
+
+from __future__ import annotations
+
+import queue
+import threading
+from functools import lru_cache
+
+import torch
+
+# A decoder retains roughly 15-25 MiB of device-side scratch space. Two
+# decoders are enough to overlap JPEG decode without mirroring the much larger
+# I/O thread count into HBM usage.
+_DECODER_POOL_SIZE = 2
+_DECODER_OPTIONS = ":num_cuda_streams=1 :fancy_upsampling=1"
+
+
+class _NvJpegDecoderPool:
+    def __init__(self, device_id: int):
+        from nvidia import nvimgcodec
+
+        self._nvimgcodec = nvimgcodec
+        self._device_id = device_id
+        self._decode_params = nvimgcodec.DecodeParams(
+            sample_format=nvimgcodec.SampleFormat.P_RGB,
+            apply_exif_orientation=False,
+        )
+        self._decoders = queue.LifoQueue(maxsize=_DECODER_POOL_SIZE)
+        self._created = 0
+        self._create_lock = threading.Lock()
+
+    def _acquire(self):
+        try:
+            return self._decoders.get_nowait()
+        except queue.Empty:
+            pass
+
+        with self._create_lock:
+            if self._created < _DECODER_POOL_SIZE:
+                decoder = self._nvimgcodec.Decoder(
+                    device_id=self._device_id,
+                    max_num_cpu_threads=1,
+                    options=_DECODER_OPTIONS,
+                )
+                self._created += 1
+                return decoder
+
+        return self._decoders.get()
+
+    def decode(self, image_bytes: bytes) -> torch.Tensor:
+        decoder = self._acquire()
+        try:
+            stream = torch.cuda.current_stream(self._device_id)
+            image = decoder.decode(
+                image_bytes,
+                params=self._decode_params,
+                cuda_stream=stream.cuda_stream,
+            )
+            if image is None:
+                raise RuntimeError("nvImageCodec could not decode the JPEG image")
+            return torch.from_dlpack(image.to_dlpack(cuda_stream=stream.cuda_stream))
+        finally:
+            self._decoders.put(decoder)
+
+
+@lru_cache(maxsize=None)
+def _get_decoder_pool(device_id: int) -> _NvJpegDecoderPool:
+    return _NvJpegDecoderPool(device_id)
+
+
+def decode_jpeg_with_fancy_upsampling(image_bytes: bytes) -> torch.Tensor:
+    """Decode a JPEG to contiguous CHW RGB uint8 on the current CUDA device.
+
+    torchvision's CUDA JPEG decoder creates nvJPEG with its default flags,
+    which use nearest-neighbor chroma upsampling. nvImageCodec exposes nvJPEG's
+    interpolated ("fancy") upsampling and exports the result to PyTorch through
+    DLPack without copying it.
+    """
+    device_id = torch.cuda.current_device()
+    image = _get_decoder_pool(device_id).decode(image_bytes)
+    if image.ndim != 3 or image.shape[0] != 3 or image.dtype != torch.uint8:
+        raise RuntimeError(
+            "nvImageCodec returned an invalid JPEG tensor: "
+            f"shape={tuple(image.shape)}, dtype={image.dtype}"
+        )
+    return image