@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ce7391712b

[Bugfix][Security] Bound embedding densification before to_dense() (#54632)

Signed-off-by: lzhan011 <35493221+lzhan011@users.noreply.github.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

lzhan011 committed Sep 1, 2026 at 03:45 UTC ce7391712b59f55a8d1254468605c5a3592d2d50

8 files changed +226 -18

docs/usage/security.md

@@ -98,6 +98,7 @@ denial of service:

98 | `VLLM_MAX_AUDIO_CLIP_FILESIZE_MB` | `25` | Maximum compressed filesize in MB for a single audio file. Enforced on all audio inputs (multimodal chat URLs, speech-to-text uploads, data: URLs, and local file paths) before decoding begins. |

99 | `VLLM_MAX_AUDIO_DECODE_DURATION_S` | `600` | Maximum decoded audio duration in seconds. Prevents compressed audio from expanding into gigabytes of float32 PCM. |

100 | `VLLM_MAX_AUDIO_DECODE_BYTES` | `268435456` (256 MiB) | Maximum float32 PCM bytes that audio decoding may allocate. Guards against sample-rate forgery where an inflated header sample rate bypasses the duration guard while the actual frame count causes a multi-GiB allocation. |

101 +| `VLLM_MAX_EMBED_DECODE_BYTES` | `2147483648` (2 GiB) | Maximum bytes a client-supplied embedding payload (`prompt_embeds`, `image_embeds`, `audio_embeds`, `video_embeds`) may allocate once densified. A sparse tensor carries its own declared shape, so a payload of a few hundred bytes can expand into hundreds of GiB. Checked before `to_dense()`, so the memory is never allocated. Set to `0` to disable. |

102

103 Setting any of these to `0` disables the corresponding limit. This is **not

104 recommended** for deployments exposed to untrusted users, as it removes the

tests/renderers/test_sparse_tensor_validation.py

+143 -1

@@ -2,7 +2,13 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """

4 Tests verify that malicious sparse tensors are rejected before they can trigger

5 -out-of-bounds memory writes during to_dense() operations.

5 +out-of-bounds memory writes, or an unbounded allocation, during to_dense()

6 +operations.

7 +

8 +A sparse tensor carries its own declared shape, so the two failure modes are

9 +independent: invalid *indices* are what the invariant checks catch, while a

10 +valid tensor with an enormous declared *shape* passes those checks and is

11 +bounded by `VLLM_MAX_EMBED_DECODE_BYTES` instead.

12 """

14 import io

@@ -12,9 +18,12 @@ import pybase64 as base64

18 import pytest

19 import torch

21 +import vllm.envs as envs

22 from vllm.exceptions import VLLMValidationError

23 from vllm.multimodal.media import AudioEmbeddingMediaIO, ImageEmbeddingMediaIO

24 +from vllm.multimodal.media.video import VideoEmbeddingMediaIO

25 from vllm.renderers.embed_utils import safe_load_prompt_embeds

26 +from vllm.utils.sparse_utils import safe_to_dense

29 @pytest.fixture

@@ -397,3 +406,136 @@ class TestSparseTensorValidationIntegration:

406

407 with pytest.raises((RuntimeError, ValueError)):

408 io_handler.load_base64("", attack_payload.decode("utf-8"))

409 +

410 +

411 +# The real attack payload: 4 TiB of float32 once densified, from a single

412 +# stored element. Only ever measured, never decoded.

413 +BOMB_SHAPE = (2**20, 2**20)

414 +

415 +# What the size-limit tests below actually build. Each of them also lowers

416 +# `VLLM_MAX_EMBED_DECODE_BYTES` to `TEST_LIMIT_BYTES`, so the guard still has

417 +# something to reject while nothing here can allocate more than 4 MiB even if

418 +# the guard regressed. Keeping the 4 TiB shape in these tests would make the

419 +# outcome depend on the runner: with `vm.overcommit_memory=0` the allocator

420 +# refuses it outright, but with `vm.overcommit_memory=1` the malloc succeeds

421 +# and the zero-fill is what fails, which would OOM the test process.

422 +OVERSIZED_SHAPE = (1024, 1024) # 4 MiB dense

423 +TEST_LIMIT_BYTES = 4096

424 +

425 +EMBEDDING_MEDIA_IO = [

426 + (ImageEmbeddingMediaIO, "image_embeds"),

427 + (AudioEmbeddingMediaIO, "audio_embeds"),

428 + (VideoEmbeddingMediaIO, "video_embeds"),

429 +]

430 +

431 +

432 +def _create_oversized_sparse_tensor(shape: tuple[int, int]) -> torch.Tensor:

433 + """A *valid* sparse tensor with one stored element and a declared `shape`.

434 +

435 + Indices are in bounds, so `check_sparse_tensor_invariants()` accepts it.

436 + That is the point: the invariant checks are not what stops this one.

437 + """

438 + indices = torch.tensor([[0], [0]])

439 + values = torch.tensor([1.0])

440 + return torch.sparse_coo_tensor(indices, values, shape, dtype=torch.float32)

441 +

442 +

443 +@pytest.fixture

444 +def small_decode_limit(monkeypatch: pytest.MonkeyPatch) -> None:

445 + monkeypatch.setenv("VLLM_MAX_EMBED_DECODE_BYTES", str(TEST_LIMIT_BYTES))

446 +

447 +

448 +class TestEmbeddingDecodeSizeLimit:

449 + """Bound the dense allocation a payload can declare.

450 +

451 + `weights_only=True` stops code execution and the invariant checks stop

452 + out-of-bounds indices, but neither bounds the declared shape.

453 + """

454 +

455 + def test_oversized_sparse_tensor_rejected(self, small_decode_limit):

456 + with pytest.raises(VLLMValidationError) as exc_info:

457 + safe_to_dense(

458 + _create_oversized_sparse_tensor(OVERSIZED_SHAPE),

459 + parameter="prompt_embeds",

460 + )

461 +

462 + assert "VLLM_MAX_EMBED_DECODE_BYTES" in str(exc_info.value)

463 + assert "prompt_embeds" in str(exc_info.value)

464 +

465 + def test_default_limit_would_reject_the_real_payload(self):

466 + """The shipped default has to bound the actual attack, not just the

467 + shrunken shape used above. Checked arithmetically so no test ever asks

468 + the allocator for 4 TiB."""

469 + dense_bytes = BOMB_SHAPE[0] * BOMB_SHAPE[1] * 4

470 + assert dense_bytes > envs.VLLM_MAX_EMBED_DECODE_BYTES > 0

471 +

472 + def test_attack_payload_is_tiny(self):

473 + """The payload really is small - this is amplification, not an upload."""

474 + assert len(_encode_tensor(_create_oversized_sparse_tensor(BOMB_SHAPE))) < 4096

475 +

476 + def test_valid_tensors_still_densify(self):

477 + assert safe_to_dense(

478 + _create_valid_dense_tensor(), parameter="prompt_embeds"

479 + ).shape == (10, 768)

480 +

481 + dense = safe_to_dense(_create_valid_sparse_tensor(), parameter="image_embeds")

482 + assert dense.shape == (3, 3)

483 + assert not dense.is_sparse

484 +

485 + def test_limit_is_configurable(self, monkeypatch: pytest.MonkeyPatch):

486 + # 32 x 32 float32 == 4096 bytes, exactly at the limit.

487 + monkeypatch.setenv("VLLM_MAX_EMBED_DECODE_BYTES", str(TEST_LIMIT_BYTES))

488 +

489 + at_limit = _create_oversized_sparse_tensor((32, 32))

490 + assert safe_to_dense(at_limit, parameter="image_embeds").shape == (32, 32)

491 +

492 + over_limit = _create_oversized_sparse_tensor((33, 32))

493 + with pytest.raises(VLLMValidationError):

494 + safe_to_dense(over_limit, parameter="image_embeds")

495 +

496 + def test_limit_can_be_disabled(self, monkeypatch: pytest.MonkeyPatch):

497 + monkeypatch.setenv("VLLM_MAX_EMBED_DECODE_BYTES", "0")

498 + tensor = _create_oversized_sparse_tensor(OVERSIZED_SHAPE)

499 +

500 + assert safe_to_dense(tensor, parameter="image_embeds").shape == OVERSIZED_SHAPE

501 +

502 + @pytest.mark.parametrize("io_cls,parameter", EMBEDDING_MEDIA_IO)

503 + def test_media_io_rejects_oversized_base64(

504 + self, io_cls, parameter, small_decode_limit

505 + ):

506 + encoded = _encode_tensor(_create_oversized_sparse_tensor(OVERSIZED_SHAPE))

507 +

508 + with pytest.raises(VLLMValidationError) as exc_info:

509 + io_cls().load_base64("", encoded.decode("utf-8"))

510 +

511 + assert parameter in str(exc_info.value)

512 +

513 + @pytest.mark.parametrize("io_cls,parameter", EMBEDDING_MEDIA_IO)

514 + def test_media_io_rejects_oversized_file(

515 + self, io_cls, parameter, tmp_path, small_decode_limit

516 + ):

517 + path = tmp_path / "embeds.pt"

518 + torch.save(_create_oversized_sparse_tensor(OVERSIZED_SHAPE), path)

519 +

520 + with pytest.raises(VLLMValidationError):

521 + io_cls().load_file(path)

522 +

523 + @pytest.mark.parametrize("io_cls,parameter", EMBEDDING_MEDIA_IO)

524 + def test_media_io_rejects_non_tensor_payload(self, io_cls, parameter):

525 + """These three had no `isinstance` check, so a non-tensor payload

526 + surfaced as an `AttributeError` (500) rather than a 4xx."""

527 + buffer = io.BytesIO()

528 + torch.save({"not": "a tensor"}, buffer)

529 + encoded = base64.b64encode(buffer.getvalue()).decode("utf-8")

530 +

531 + with pytest.raises(VLLMValidationError) as exc_info:

532 + io_cls().load_base64("", encoded)

533 +

534 + assert "torch.Tensor" in str(exc_info.value)

535 +

536 + @pytest.mark.parametrize("io_cls,parameter", EMBEDDING_MEDIA_IO)

537 + def test_media_io_still_accepts_valid_embeddings(self, io_cls, parameter):

538 + tensor = torch.randn(4, 16, dtype=torch.float32)

539 + encoded = _encode_tensor(tensor)

540 +

541 + assert torch.equal(io_cls().load_base64("", encoded.decode("utf-8")), tensor)

vllm/envs.py

+12

@@ -82,6 +82,7 @@ if TYPE_CHECKING:

82 VLLM_MAX_AUDIO_DECODE_DURATION_S: int = 600

83 VLLM_MAX_AUDIO_DECODE_BYTES: int = 268_435_456

84 VLLM_MAX_AUDIO_PREPROCESS_WORKERS: int = max(1, min(os.cpu_count() or 1, 2))

85 + VLLM_MAX_EMBED_DECODE_BYTES: int = 2_147_483_648

86 VLLM_MAX_IMAGE_PIXELS: int = 178_956_970

87 VLLM_VIDEO_LOADER_BACKEND: str = "opencv"

88 VLLM_MEDIA_CONNECTOR: str = "http"

@@ -1011,6 +1012,16 @@ environment_variables: dict[str, Callable[[], Any]] = {

1012 "VLLM_MAX_AUDIO_DECODE_BYTES": lambda: int(

1013 os.getenv("VLLM_MAX_AUDIO_DECODE_BYTES", "268435456")

1014 ),

1015 + # Maximum bytes a client-supplied embedding payload (`prompt_embeds`,

1016 + # `image_embeds`, `audio_embeds`, `video_embeds`) may allocate once it is

1017 + # densified. A sparse tensor carries its own declared shape, so a payload

1018 + # of a few hundred bytes can expand into hundreds of GiB. The limit is

1019 + # checked before `to_dense()` so the memory is never allocated. Set to 0

1020 + # to disable. Default is 2 GiB, which covers a 128K-token float32

1021 + # embedding at hidden_size 4096.

1022 + "VLLM_MAX_EMBED_DECODE_BYTES": lambda: int(

1023 + os.getenv("VLLM_MAX_EMBED_DECODE_BYTES", "2147483648")

1024 + ),

1025 # Maximum number of worker threads used for STT preprocessing. The default

1026 # intentionally caps at 2 because that performed best in profiling.

1027 # https://github.com/vllm-project/vllm/pull/44612#issuecomment-4662757781

@@ -2325,6 +2336,7 @@ def compile_factors() -> dict[str, object]:

2336 "VLLM_MAX_AUDIO_DECODE_DURATION_S",

2337 "VLLM_MAX_AUDIO_DECODE_BYTES",

2338 "VLLM_MAX_AUDIO_PREPROCESS_WORKERS",

2339 + "VLLM_MAX_EMBED_DECODE_BYTES",

2340 "VLLM_MAX_IMAGE_PIXELS",

2341 "VLLM_VIDEO_LOADER_BACKEND",

2342 "VLLM_MEDIA_CONNECTOR",

vllm/multimodal/media/audio.py

+6 -3

@@ -16,7 +16,10 @@ from vllm.multimodal.audio import resample_audio_pyav

16 from vllm.utils.import_utils import PlaceholderModule

17 from vllm.utils.mem_constants import MiB_bytes

18 from vllm.utils.serial_utils import tensor2base64

19 -from vllm.utils.sparse_utils import check_sparse_tensor_invariants_threadsafe

19 +from vllm.utils.sparse_utils import (

20 + check_sparse_tensor_invariants_threadsafe,

21 + safe_to_dense,

22 +)

24 from .base import MediaIO

@@ -348,7 +351,7 @@ class AudioEmbeddingMediaIO(MediaIO[torch.Tensor]):

351 buffer = BytesIO(data)

352 with check_sparse_tensor_invariants_threadsafe():

353 tensor = torch.load(buffer, weights_only=True)

351 - return tensor.to_dense()

354 + return safe_to_dense(tensor, parameter="audio_embeds")

355

356 def load_base64(self, media_type: str, data: str) -> torch.Tensor:

357 return self.load_bytes(pybase64.b64decode(data, validate=True))

@@ -356,7 +359,7 @@ class AudioEmbeddingMediaIO(MediaIO[torch.Tensor]):

359 def load_file(self, filepath: Path) -> torch.Tensor:

360 with check_sparse_tensor_invariants_threadsafe():

361 tensor = torch.load(filepath, weights_only=True)

359 - return tensor.to_dense()

362 + return safe_to_dense(tensor, parameter="audio_embeds")

363

364 def encode_base64(self, media: torch.Tensor) -> str:

365 return tensor2base64(media)

vllm/multimodal/media/image.py

+6 -3

@@ -11,7 +11,10 @@ from PIL import Image

12 import vllm.envs as envs

13 from vllm.utils.serial_utils import tensor2base64

14 -from vllm.utils.sparse_utils import check_sparse_tensor_invariants_threadsafe

14 +from vllm.utils.sparse_utils import (

15 + check_sparse_tensor_invariants_threadsafe,

16 + safe_to_dense,

17 +)

19 from ..image import convert_image_mode, normalize_image, rgba_to_rgb

20 from .base import MediaIO, MediaWithBytes

@@ -135,7 +138,7 @@ class ImageEmbeddingMediaIO(MediaIO[torch.Tensor]):

138 buffer = BytesIO(data)

139 with check_sparse_tensor_invariants_threadsafe():

140 tensor = torch.load(buffer, weights_only=True)

138 - return tensor.to_dense()

141 + return safe_to_dense(tensor, parameter="image_embeds")

142

143 def _load_numpy(self, data: bytes) -> torch.Tensor:

144 with BytesIO(data) as buffer:

@@ -156,7 +159,7 @@ class ImageEmbeddingMediaIO(MediaIO[torch.Tensor]):

159

160 with check_sparse_tensor_invariants_threadsafe():

161 tensor = torch.load(filepath, weights_only=True)

159 - return tensor.to_dense()

162 + return safe_to_dense(tensor, parameter="image_embeds")

163

164 def encode_base64(self, media: torch.Tensor) -> str:

165 return tensor2base64(media)

vllm/multimodal/media/video.py

+6 -3

@@ -14,7 +14,10 @@ from PIL import Image

14 from vllm import envs

15 from vllm.logger import init_logger

16 from vllm.utils.serial_utils import tensor2base64

17 -from vllm.utils.sparse_utils import check_sparse_tensor_invariants_threadsafe

17 +from vllm.utils.sparse_utils import (

18 + check_sparse_tensor_invariants_threadsafe,

19 + safe_to_dense,

20 +)

22 from ..video import VIDEO_LOADER_REGISTRY

23 from .base import MediaIO, MediaWithBytes

@@ -215,7 +218,7 @@ class VideoEmbeddingMediaIO(MediaIO[torch.Tensor]):

218 buffer = BytesIO(data)

219 with check_sparse_tensor_invariants_threadsafe():

220 tensor = torch.load(buffer, weights_only=True)

218 - return tensor.to_dense()

221 + return safe_to_dense(tensor, parameter="video_embeds")

222

223 def _load_numpy(self, data: bytes) -> torch.Tensor:

224 with BytesIO(data) as buffer:

@@ -236,7 +239,7 @@ class VideoEmbeddingMediaIO(MediaIO[torch.Tensor]):

239

240 with check_sparse_tensor_invariants_threadsafe():

241 tensor = torch.load(filepath, weights_only=True)

239 - return tensor.to_dense()

242 + return safe_to_dense(tensor, parameter="video_embeds")

243

244 def encode_base64(self, media: torch.Tensor) -> str:

245 return tensor2base64(media)

vllm/renderers/embed_utils.py

+5 -7

@@ -8,7 +8,10 @@ import torch

9 from vllm.exceptions import VLLMValidationError

10 from vllm.utils.async_utils import make_async

11 -from vllm.utils.sparse_utils import check_sparse_tensor_invariants_threadsafe

11 +from vllm.utils.sparse_utils import (

12 + check_sparse_tensor_invariants_threadsafe,

13 + safe_to_dense,

14 +)

16 if TYPE_CHECKING:

17 from vllm.config import ModelConfig

@@ -30,12 +33,7 @@ def safe_load_prompt_embeds(

33 weights_only=True,

34 map_location=torch.device("cpu"),

35 )

33 - if not isinstance(tensor, torch.Tensor):

34 - raise VLLMValidationError(

35 - "`prompt_embeds` payload did not deserialize to a torch.Tensor.",

36 - parameter="prompt_embeds",

37 - )

38 - tensor = tensor.to_dense()

36 + tensor = safe_to_dense(tensor, parameter="prompt_embeds")

38 if tensor.dim() > 2:

39 tensor = tensor.squeeze(0)

vllm/utils/sparse_utils.py

+47 -1

@@ -1,7 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """

4 -Thread-safe wrapper for sparse tensor invariant validation.

4 +Safe deserialization helpers for client-supplied embedding tensors.

5 +

6 +`check_sparse_tensor_invariants_threadsafe()` is a thread-safe wrapper for

7 +sparse tensor invariant validation. `safe_to_dense()` bounds the memory a

8 +deserialized tensor may allocate when densified.

10 PyTorch's `torch.sparse.check_sparse_tensor_invariants()` context manager

11 manipulates a **process-global** flag (save/enable/restore). When multiple

@@ -18,6 +22,10 @@ import threading

23 import torch

25 +import vllm.envs as envs

26 +from vllm.exceptions import VLLMValidationError

27 +from vllm.utils.mem_constants import MiB_bytes

28 +

29 _SPARSE_LOAD_LOCK = threading.Lock()

@@ -25,3 +33,41 @@ _SPARSE_LOAD_LOCK = threading.Lock()

33 def check_sparse_tensor_invariants_threadsafe():

34 with _SPARSE_LOAD_LOCK, torch.sparse.check_sparse_tensor_invariants():

35 yield

36 +

37 +

38 +def safe_to_dense(tensor: object, *, parameter: str) -> torch.Tensor:

39 + """Densify a client-supplied embedding tensor within a memory budget.

40 +

41 + `torch.load(..., weights_only=True)` prevents arbitrary code execution but

42 + not memory amplification: a sparse COO tensor carries its own declared

43 + shape, so the dense allocation is bounded by that shape rather than by the

44 + request. A 2.6 KB payload declaring `(30000, 30000)` materializes 3.4 GiB,

45 + and the shape can be picked to fit whatever host is being targeted. The

46 + invariant checks only reject *invalid* tensors (indices outside the

47 + declared shape); a single non-zero element inside a `(2**20, 2**20)` shape

48 + is perfectly valid and passes them.

49 +

50 + So reject oversized payloads based on the dense size they declare, before

51 + `to_dense()` allocates anything.

52 + """

53 + if not isinstance(tensor, torch.Tensor):

54 + raise VLLMValidationError(

55 + f"`{parameter}` payload did not deserialize to a torch.Tensor.",

56 + parameter=parameter,

57 + )

58 +

59 + max_bytes = envs.VLLM_MAX_EMBED_DECODE_BYTES

60 + if max_bytes > 0:

61 + dense_bytes = tensor.numel() * tensor.element_size()

62 + if dense_bytes > max_bytes:

63 + raise VLLMValidationError(

64 + f"`{parameter}` declares shape {tuple(tensor.shape)}, which "

65 + f"would allocate {dense_bytes} bytes "

66 + f"({dense_bytes / MiB_bytes:.0f} MiB) once densified, "

67 + f"exceeding the {max_bytes} byte limit. Set "

68 + f"VLLM_MAX_EMBED_DECODE_BYTES to increase this limit.",

69 + parameter=parameter,

70 + value=dense_bytes,

71 + )

72 +

73 + return tensor.to_dense()