xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-04-12 06:37:03 +08:00
Code Issues Packages Projects Releases Wiki Activity

Use math.prod instead of np.prod for trivial ops (#14142)

Browse Source
This commit is contained in:
Zhanwen Chen 2025-03-04 00:09:22 -05:00 committed by GitHub
parent bf13d40972
commit 66233af7b6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
vllm/worker/cache_engine.py
View File

@ -1,8 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
"""CacheEngine class for managing the KV cache."""
from math import prod
from typing import List
import numpy as np
import torch
from vllm import envs
@ -90,7 +90,7 @@ class CacheEngine:
# NOTE this assumption currently only holds for MLA so we only apply
# this optimization when `use_mla` is true
entry_shape = kv_cache_shape[2:]
entry_size = np.prod(entry_shape)
entry_size = prod(entry_shape)
alloc_entry_size = align_to_256bytes(entry_size, self.dtype)
alloc_shape = (*kv_cache_shape[:2], alloc_entry_size)
else: