wehub-resource-sync
65 lines
2.5 KiB
C++
65 lines
2.5 KiB
C++
// SPDX-License-Identifier: Apache-2.0
|
|
#pragma once
|
|
|
|
#include <cstddef>
|
|
#include <cstdint>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#include "bitmap.h"
|
|
|
|
namespace lmcache {
|
|
|
|
namespace storage_manager {
|
|
|
|
/**
|
|
* @brief Fold per-(group, chunk, rank) presence into servable prefix lengths.
|
|
*
|
|
* For each object group, computes which prefix lengths it can serve under its
|
|
* rule (full attention or a cross-chunk sliding window), and intersects across
|
|
* groups. The result feeds :func:`Bitmap::highest_set_bit`; the model-wide hit
|
|
* length is that index plus one (``-1`` -> hit length 0), then :func:`unfold`.
|
|
*
|
|
* The input ``found`` is group-major / chunk-major / rank-minor: bit
|
|
* ``g * (num_chunks * num_ranks) + j * num_ranks + r`` is set iff chunk ``j``
|
|
* of object group ``g`` is present for kv_rank ``r``. A chunk counts as present
|
|
* for a group only when every kv_rank shard is present.
|
|
*
|
|
* @param found Presence bitmap of length
|
|
* ``group_windows.size() * num_chunks * num_ranks``.
|
|
* @param num_chunks Number of LMCache chunks in the request.
|
|
* @param num_ranks Number of kv_rank shards per chunk.
|
|
* @param group_windows Per-object-group cross-chunk sliding-window size in
|
|
* chunks, in object-group order; ``<= 0`` means full attention.
|
|
*
|
|
* @return A bitmap of size ``num_chunks``; bit ``j`` set iff every group can
|
|
* serve a length-``j + 1`` prefix.
|
|
*/
|
|
Bitmap fold(const Bitmap& found, size_t num_chunks, size_t num_ranks,
|
|
const std::vector<int64_t>& group_windows);
|
|
|
|
/**
|
|
* @brief Expand a model-wide hit length into the per-group retain mask.
|
|
*
|
|
* Each group retains the chunks it needs to serve ``hit_length``: ``[0,
|
|
* hit_length)`` for full attention, ``[hit_length - window, hit_length)`` for a
|
|
* sliding window. The mask is over the same ranked layout as :func:`fold`'s
|
|
* input (all kv_ranks of each retained ``(group, chunk)`` set).
|
|
*
|
|
* @param hit_length Model-wide prefix hit length in chunks (clamped to
|
|
* ``num_chunks``).
|
|
* @param num_chunks Number of LMCache chunks in the request.
|
|
* @param num_ranks Number of kv_rank shards per chunk.
|
|
* @param group_windows Per-object-group cross-chunk sliding-window size in
|
|
* chunks, in object-group order; ``<= 0`` means full attention.
|
|
*
|
|
* @return Retain mask of length
|
|
* ``group_windows.size() * num_chunks * num_ranks``.
|
|
*/
|
|
Bitmap unfold(size_t hit_length, size_t num_chunks, size_t num_ranks,
|
|
const std::vector<int64_t>& group_windows);
|
|
|
|
} // namespace storage_manager
|
|
|
|
} // namespace lmcache
|