A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 21:00 UTC
Use math.prod instead of np.prod for trivial ops (#14142)
Zhanwen Chen committed
Mar 4, 2025 at 00:09 UTC
66233af7b6e4217653f1a9952180d68376af7d2a
1 file changed
+2
-2
vllm/worker/cache_engine.py
+2
-2
@@ -1,8 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
"""CacheEngine class for managing the KV cache."""
3
+from math import prod
4
from typing import List
5
5
-import numpy as np
6
import torch
7
8
from vllm import envs
@@ -90,7 +90,7 @@ class CacheEngine:
90
# NOTE this assumption currently only holds for MLA so we only apply
91
# this optimization when `use_mla` is true
92
entry_shape = kv_cache_shape[2:]
93
- entry_size = np.prod(entry_shape)
93
+ entry_size = prod(entry_shape)
94
alloc_entry_size = align_to_256bytes(entry_size, self.dtype)
95
alloc_shape = (*kv_cache_shape[:2], alloc_entry_size)
96
else: