[Feature]: Use math.prod instead of np.prod for trivial ops

[zhanwenchen]

🚀 The feature, motivation and pitch

Currently, https://github.com/vllm-project/vllm/blob/main/vllm/worker/cache_engine.py#L93 imports numpy solely for the purpose of entry_size = np.prod(entry_shape), where entry_shape = kv_cache_shape[2:] is an iterable of trivial size. Using the native Python 3.8 math.prod (https://docs.python.org/3/library/math.html#math.prod) would make the both imports and compute overhead lighter.

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[hmellor]

Interestingly, when microbenchmarking these two ops, math.prod appears to be slightly faster for this use case too

import math
import numpy as np
import timeit

# Test data: list of 2D tensor shapes
shapes = [
    (10, 20),
    (32, 64),
    (128, 256),
    (512, 1024),
    (2048, 4096)
]

# Number of iterations
iterations = 100000

# Benchmark math.prod
def test_math_prod():
    for shape in shapes:
        _ = math.prod(shape)

math_time = timeit.timeit(test_math_prod, number=iterations)

# Benchmark np.prod
def test_numpy_prod():
    for shape in shapes:
        _ = np.prod(shape)

numpy_time = timeit.timeit(test_numpy_prod, number=iterations)

# Print results
print(f"Testing with {iterations} iterations on {len(shapes)} different shapes:")
print(f"math.prod time: {math_time:.6f} seconds")
print(f"np.prod time:   {numpy_time:.6f} seconds")

if math_time < numpy_time:
    print(f"math.prod is {numpy_time/math_time:.2f}x faster than np.prod")
else:
    print(f"np.prod is {math_time/numpy_time:.2f}x faster than math.prod")