wehub-resource-sync
95 lines
3.3 KiB
Python
95 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
"""Microbenchmark for the ranked fold (:func:`fold_unfold_ranked`).
|
|
|
|
Compares the pure-Python reference against the native C++ implementation across
|
|
request sizes, including a DeepSeek-scale hybrid case (1M tokens, 8 object
|
|
groups mixing full attention and sliding window). Run with::
|
|
|
|
python benchmarks/microbenchmark/bitmap_ops_benchmark.py
|
|
|
|
The native op scans the packed ``Bitmap`` buffer directly -- no Python per-bit
|
|
loop and no ``Bitmap``<->tensor conversion -- so it stays sub-millisecond even
|
|
at multi-million-key scale where the Python scan takes hundreds of ms.
|
|
"""
|
|
|
|
# Standard
|
|
from collections.abc import Sequence
|
|
import time
|
|
|
|
# First Party
|
|
from lmcache.native_storage_ops import Bitmap
|
|
from lmcache.v1.distributed.bitmap_ops import fold_unfold_ranked, highest_set_bit
|
|
from lmcache.v1.distributed.bitmap_ops.fold import _fold_python, _unfold_python
|
|
|
|
|
|
def _python_pipeline(found, num_chunks, num_ranks, group_windows):
|
|
"""Pure-Python fold -> highest_set_bit -> unfold (no native ops)."""
|
|
servable = _fold_python(found, num_chunks, num_ranks, group_windows)
|
|
hit = highest_set_bit(servable) + 1 # -1 (no servable prefix) -> 0
|
|
return hit, _unfold_python(hit, num_chunks, num_ranks, group_windows)
|
|
|
|
|
|
def _best_ms(fn, reps: int) -> float:
|
|
"""Best wall-clock time of ``fn`` over ``reps`` runs, in milliseconds."""
|
|
best = float("inf")
|
|
for _ in range(reps):
|
|
start = time.perf_counter()
|
|
fn()
|
|
best = min(best, time.perf_counter() - start)
|
|
return best * 1e3
|
|
|
|
|
|
def bench_case(
|
|
label: str,
|
|
num_chunks: int,
|
|
num_ranks: int,
|
|
group_windows: Sequence[int],
|
|
present_fraction: float,
|
|
reps: int = 5,
|
|
) -> None:
|
|
"""Print Python vs native timings for one (size, fill) configuration."""
|
|
num_keys = len(group_windows) * num_chunks * num_ranks
|
|
if present_fraction >= 1.0:
|
|
found = Bitmap(num_keys, num_keys)
|
|
else:
|
|
found = Bitmap(num_keys, int(num_keys * present_fraction))
|
|
|
|
windows = list(group_windows)
|
|
py_ms = _best_ms(
|
|
lambda: _python_pipeline(found, num_chunks, num_ranks, windows),
|
|
reps,
|
|
)
|
|
native_ms = _best_ms(
|
|
lambda: fold_unfold_ranked(found, num_chunks, num_ranks, windows),
|
|
reps,
|
|
)
|
|
speedup = py_ms / native_ms if native_ms else float("inf")
|
|
print(
|
|
f"{label:<34}keys={num_keys:>9} python={py_ms:>9.2f}ms "
|
|
f"native={native_ms:>8.3f}ms speedup={speedup:>7.1f}x"
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
"""Run the benchmark grid."""
|
|
# 8 groups, mix of full attention and sliding window (DeepSeek-like hybrid).
|
|
dpsk_windows = (-1, -1, -1, -1, 4, 4, 8, 1)
|
|
|
|
print("== DeepSeek 1M tokens @ chunk_size=256 (num_chunks=4096), all present ==")
|
|
bench_case("dpsk, world_size=1", 4096, 1, dpsk_windows, 1.0)
|
|
bench_case("dpsk, world_size=8", 4096, 8, dpsk_windows, 1.0)
|
|
|
|
print("\n== same, 50% prefix present (realistic) ==")
|
|
bench_case("dpsk, world_size=8", 4096, 8, dpsk_windows, 0.5)
|
|
|
|
print("\n== small request ==")
|
|
bench_case("4K tokens @256 (16 chunks)", 16, 8, dpsk_windows, 1.0)
|
|
|
|
print("\n== stress: chunk_size=16 -> 62500 chunks (4M keys) ==")
|
|
bench_case("stress, world_size=8", 62500, 8, dpsk_windows, 1.0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|