1
2
3
4
5
6

# vLLM的Block结构
class Block:
    block_id: int           # 块编号
    block_size: int         # 块大小（tokens数）
    ref_count: int          # 引用计数
    physical_block: int     # 物理块位置

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

class LLMEngine:
    def __init__(self, ...):
        self.scheduler = Scheduler(...)      # 调度器
        self.model_executor = ModelExecutor(...)  # 执行器
        self.tokenizer = Tokenizer(...)      # 分词器

    def add_request(self, request_id, prompt, ...):
        """添加推理请求"""
        pass

    def step(self):
        """执行一步推理"""
        # 1. 调度器选择要处理的请求
        scheduler_outputs = self.scheduler.schedule()

        # 2. 执行模型前向传播
        outputs = self.model_executor.execute(scheduler_outputs)

        # 3. 处理输出，更新状态
        return self._process_outputs(outputs)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

class Scheduler:
    def __init__(self, ...):
        self.running: List[SequenceGroup] = []    # 正在运行的请求
        self.waiting: deque[SequenceGroup] = deque()  # 等待中的请求
        self.swapped: List[SequenceGroup] = []    # 被换出的请求
        self.block_manager = BlockManager(...)    # 块管理器

    def schedule(self) -> SchedulerOutputs:
        """调度决策"""
        # 1. 检查running请求是否可以继续
        running_queue = self._schedule_running()

        # 2. 尝试恢复swapped请求
        swapped_in = self._schedule_swapped()

        # 3. 尝试调度waiting请求
        prefills = self._schedule_prefills()

        return SchedulerOutputs(running_queue, swapped_in, prefills)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

class ModelRunner:
    def __init__(self, ...):
        self.model = load_model(...)
        self.kv_cache = KVCache(...)

    def execute_model(self, scheduler_output):
        """执行模型前向传播"""
        # 1. 准备输入
        input_ids, positions = self._prepare_inputs(scheduler_output)

        # 2. 前向传播
        with set_forward_context(self.attn_metadata):
            hidden_states = self.model.forward(input_ids, positions)

        # 3. 采样输出token
        next_tokens = self.sampler(hidden_states)

        return next_tokens

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

# 单个Block的内存布局
class KVCacheBlock:
    """
    KV Cache Block 内存布局:

    K Cache:
    ┌─────────────────────────────────────────┐
    │ Token 0: [head0][head1]...[headN-1]     │
    │ Token 1: [head0][head1]...[headN-1]     │
    │ ...                                     │
    │ Token B-1: [head0][head1]...[headN-1]   │
    └─────────────────────────────────────────┘

    V Cache: 相同布局

    其中:
    - B = block_size (例如16 tokens)
    - N = num_kv_heads (例如8)
    - 每个head的维度 = head_dim (例如128)
    """

    def __init__(self, block_size, num_kv_heads, head_dim, dtype):
        self.block_size = block_size
        self.num_kv_heads = num_kv_heads
        self.head_dim = head_dim

        # K和V各占用的字节数
        single_size = block_size * num_kv_heads * head_dim * dtype.itemsize
        self.k_cache = torch.empty(single_size, dtype=dtype)
        self.v_cache = torch.empty(single_size, dtype=dtype)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46

class BlockManager:
    def __init__(self, num_gpu_blocks, num_cpu_blocks, block_size):
        # GPU块池
        self.gpu_allocator = BlockAllocator(
            device="cuda",
            num_blocks=num_gpu_blocks,
            block_size=block_size
        )

        # CPU块池（用于swap）
        self.cpu_allocator = BlockAllocator(
            device="cpu",
            num_blocks=num_cpu_blocks,
            block_size=block_size
        )

        # 每个sequence的块表
        self.block_tables: Dict[int, BlockTable] = {}

    def allocate(self, seq_id: int, num_blocks: int) -> List[int]:
        """分配指定数量的物理块"""
        blocks = self.gpu_allocator.allocate(num_blocks)
        self.block_tables[seq_id].extend(blocks)
        return blocks

    def free(self, seq_id: int):
        """释放sequence的所有块"""
        blocks = self.block_tables.pop(seq_id)
        self.gpu_allocator.free(blocks)

    def swap_out(self, seq_id: int) -> Dict[int, int]:
        """将GPU块换出到CPU"""
        gpu_blocks = self.block_tables[seq_id]
        cpu_blocks = self.cpu_allocator.allocate(len(gpu_blocks))

        # 建立映射关系
        mapping = dict(zip(gpu_blocks, cpu_blocks))

        # 异步复制数据
        self._copy_blocks(gpu_blocks, cpu_blocks, "cuda", "cpu")

        # 释放GPU块
        self.gpu_allocator.free(gpu_blocks)
        self.block_tables[seq_id] = cpu_blocks

        return mapping

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

def paged_attention(
    query: torch.Tensor,          # [num_tokens, num_heads, head_dim]
    key_cache: torch.Tensor,      # [num_blocks, block_size, num_kv_heads, head_dim]
    value_cache: torch.Tensor,    # [num_blocks, block_size, num_kv_heads, head_dim]
    block_tables: torch.Tensor,   # [num_seqs, max_num_blocks]
    context_lens: torch.Tensor,   # [num_seqs]
    scale: float,
):
    """
    Paged Attention前向传播

    核心思路:
    1. 根据block_table找到每个序列的物理块
    2. 从非连续的物理块中gather KV
    3. 计算注意力权重和输出
    """
    # CUDA kernel实现，高效处理非连续访问
    output = torch.ops.vllm.paged_attention(
        query, key_cache, value_cache,
        block_tables, context_lens, scale
    )
    return output

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

class KVCache:
    def __init__(self, num_layers, num_blocks, block_size,
                 num_kv_heads, head_dim, dtype, device):
        # 形状: [num_layers, 2, num_blocks, block_size, num_kv_heads, head_dim]
        # 2 表示 K 和 V
        self.cache = torch.empty(
            (num_layers, 2, num_blocks, block_size, num_kv_heads, head_dim),
            dtype=dtype,
            device=device
        )

    def get_layer_cache(self, layer_idx):
        """获取某一层的KV Cache"""
        k_cache = self.cache[layer_idx, 0]  # [num_blocks, block_size, heads, dim]
        v_cache = self.cache[layer_idx, 1]
        return k_cache, v_cache

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

class CopyOnWriteBlock:
    """
    Copy-on-Write块管理

    场景: 多个请求共享相同的前缀
    - 请求A: [System Prompt][Query A]
    - 请求B: [System Prompt][Query B]

    共享System Prompt的KV块，当任一方需要修改时才复制
    """

    def __init__(self):
        self.ref_count = 1  # 引用计数
        self.is_cow = False  # 是否为CoW块

    def acquire(self):
        """增加引用"""
        self.ref_count += 1

    def release(self) -> bool:
        """释放引用，返回是否可以释放物理块"""
        self.ref_count -= 1
        return self.ref_count == 0

    def copy_on_write(self, allocator) -> 'CopyOnWriteBlock':
        """写时复制"""
        if self.ref_count == 1:
            # 独占，无需复制
            return self

        # 分配新块并复制数据
        new_block = allocator.allocate(1)[0]
        copy_data(self, new_block)

        # 减少原块引用
        self.release()

        return new_block

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

# UCM的Block标识
class UCMBlock:
    def __init__(self):
        # vLLM的物理块ID
        self.vllm_block_id: int = -1

        # UCM的内容哈希（MD5）
        self.content_hash: bytes = b""

        # 块内的token IDs
        self.token_ids: List[int] = []

        # 是否在外部存储中存在
        self.in_external_storage: bool = False

    def compute_hash(self, parent_hash: bytes, token_ids: List[int]) -> bytes:
        """
        计算块的内容哈希

        哈希链:
        Block0: hash(seed, tokens[0:16])
        Block1: hash(Block0_hash, tokens[16:32])
        Block2: hash(Block1_hash, tokens[32:48])
        ...
        """
        import hashlib
        content = parent_hash + bytes(token_ids)
        return hashlib.md5(content).digest()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

class RequestHasher:
    """UCM请求哈希器"""

    def __init__(self, vllm_config, rank_id):
        # 元数据作为种子
        model = vllm_config.model
        world_size = vllm_config.tensor_parallel_size
        dtype = str(vllm_config.dtype)

        meta = f"{model}:{world_size}:{dtype}:{rank_id}"
        self.meta_bytes = meta.encode("utf-8")

    def compute_block_hashes(self, token_ids: List[int],
                             block_size: int) -> List[bytes]:
        """计算所有完整块的哈希"""
        import hashlib

        # 种子哈希
        parent_hash = hashlib.md5(
            b"UCMHASHSEED" + self.meta_bytes
        ).digest()

        block_hashes = []
        for start in range(0, len(token_ids), block_size):
            end = start + block_size
            if end > len(token_ids):
                break  # 只处理完整块

            block_tokens = token_ids[start:end]
            content = parent_hash + bytes(block_tokens)
            block_hash = hashlib.md5(content).digest()

            block_hashes.append(block_hash)
            parent_hash = block_hash  # 链式传递

        return block_hashes

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

class AutomaticPrefixCaching:
    """
    vLLM的自动前缀缓存实现

    核心思路:
    1. 使用Radix Tree存储已计算的前缀
    2. 新请求匹配最长公共前缀
    3. 只计算未命中的部分
    """

    def __init__(self):
        self.radix_tree = RadixTree()

    def match_prefix(self, token_ids: List[int]) -> Tuple[int, List[int]]:
        """
        查找最长匹配前缀

        Returns:
            matched_len: 匹配的token数
            block_ids: 可复用的物理块ID
        """
        node = self.radix_tree.root
        matched_len = 0
        block_ids = []

        for i in range(0, len(token_ids), self.block_size):
            block_tokens = tuple(token_ids[i:i+self.block_size])

            if block_tokens in node.children:
                node = node.children[block_tokens]
                matched_len += self.block_size
                block_ids.append(node.block_id)
            else:
                break

        return matched_len, block_ids

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

class UCMPrefixCache:
    """UCM增强的前缀缓存"""

    def __init__(self, store: UcmKVStoreBaseV1):
        self.store = store  # 外部存储后端
        self.local_cache = {}  # 本地内存缓存

    def lookup(self, block_hashes: List[bytes]) -> List[bool]:
        """
        查询块是否存在

        先查本地缓存，再查外部存储
        """
        results = []
        external_queries = []

        for i, hash_ in enumerate(block_hashes):
            if hash_ in self.local_cache:
                results.append(True)
            else:
                results.append(None)  # 待查询
                external_queries.append((i, hash_))

        # 批量查询外部存储
        if external_queries:
            hashes = [h for _, h in external_queries]
            external_results = self.store.lookup(hashes)

            for (i, _), exists in zip(external_queries, external_results):
                results[i] = exists

        return results

    def load_blocks(self, block_hashes: List[bytes],
                   dst_blocks: List[int]) -> Task:
        """从外部存储加载块到GPU"""
        return self.store.load_data(
            block_ids=block_hashes,
            shard_index=[0] * len(block_hashes),
            dst_addr=self._compute_addresses(dst_blocks)
        )

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46

class KVConnectorBase_V1(ABC):
    """vLLM的KV连接器基类"""

    class KVConnectorRole(Enum):
        SCHEDULER = 0  # 调度器进程
        WORKER = 1     # 执行器进程

    def __init__(self, vllm_config, role):
        self.vllm_config = vllm_config
        self.role = role

    # === 调度器端方法 ===

    @abstractmethod
    def get_num_new_matched_tokens(
        self,
        request: Request,
        num_computed_tokens: int
    ) -> Tuple[int, bool]:
        """查询外部存储中有多少token可以复用"""
        pass

    @abstractmethod
    def build_connector_meta(
        self,
        scheduler_output: SchedulerOutput
    ) -> KVConnectorMetadata:
        """构建调度器到Worker的元数据"""
        pass

    # === Worker端方法 ===

    @abstractmethod
    def register_kv_caches(self, kv_caches: Dict[str, torch.Tensor]):
        """注册KV Cache张量"""
        pass

    @abstractmethod
    def start_load_kv(self, forward_context: ForwardContext):
        """开始异步加载KV"""
        pass

    @abstractmethod
    def wait_for_save(self):
        """等待保存完成"""
        pass

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

class UCMConnector(KVConnectorBase_V1):
    """UCM的主连接器"""

    def __init__(self, vllm_config, role):
        super().__init__(vllm_config, role)

        # 初始化存储后端
        self.store = self._init_store()

        # 初始化哈希器
        self.hasher = RequestHasher(vllm_config, self.rank)

        # 请求元数据
        self.request_metas: Dict[str, RequestMeta] = {}

    def get_num_new_matched_tokens(self, request, num_computed_tokens):
        # 1. 计算块哈希
        block_hashes = self.hasher.compute_block_hashes(
            request.prompt_token_ids,
            self.block_size
        )

        # 2. 查询外部存储
        lookup_results = self.store.lookup(block_hashes)

        # 3. 统计连续命中
        hit_blocks = 0
        for exists in lookup_results:
            if not exists:
                break
            hit_blocks += 1

        return hit_blocks * self.block_size, False

    def start_load_kv(self, forward_context):
        # 获取需要加载的块
        meta = self._get_connector_metadata()

        for req_id, dispatch_meta in meta.request_meta.items():
            ucm_block_ids, vllm_block_ids = dispatch_meta.load_block_ids

            if ucm_block_ids:
                # 计算目标地址
                addrs = self._compute_addresses(vllm_block_ids)

                # 异步加载
                task = self.store.load_data(ucm_block_ids, [0]*len(ucm_block_ids), addrs)
                self.pending_loads[req_id] = task

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

# integration/vllm/ucm_connector.py 核心结构

class UCMConnectorMetadata(KVConnectorMetadata):
    """调度器到Worker的元数据传递"""
    request_meta: Dict[str, RequestDispatchMeta]

class UCMDirectConnector(KVConnectorBase_V1):
    """直接连接器实现"""

    def __init__(self, vllm_config, role):
        # 存储后端初始化
        self.store = self._init_store()

        # 请求哈希器
        self.hasher = RequestHasher(vllm_config, self.rank)

        # KV Cache指针
        self.k_base_ptrs: np.ndarray = None
        self.v_base_ptrs: np.ndarray = None

    # 实现KVConnector接口的所有方法...

class UCMConnector:
    """工厂类，根据配置选择具体实现"""

    @classmethod
    def create(cls, vllm_config, role) -> KVConnectorBase_V1:
        config = vllm_config.kv_transfer_config.kv_connector_extra_config

        if "blend" in config:
            return UCMBlendConnector(vllm_config, role)
        else:
            return UCMDirectConnector(vllm_config, role)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

# store/ucmstore_v1.py 核心接口

class Task(ABC):
    """异步任务句柄"""
    pass

class UcmKVStoreBaseV1(ABC):
    """存储后端抽象基类"""

    @abstractmethod
    def cc_store(self) -> int:
        """返回C++底层指针"""

    @abstractmethod
    def lookup(self, block_ids: List[bytes]) -> List[bool]:
        """查询块是否存在"""

    @abstractmethod
    def prefetch(self, block_ids: List[bytes]) -> None:
        """预取块到缓存"""

    @abstractmethod
    def load(self, block_ids, shard_index, dst_tensor) -> Task:
        """加载KV到设备张量"""

    @abstractmethod
    def dump(self, block_ids, shard_index, src_tensor) -> Task:
        """保存KV到存储"""

    @abstractmethod
    def load_data(self, block_ids, shard_index, dst_addr) -> Task:
        """低级加载到设备地址"""

    @abstractmethod
    def dump_data(self, block_ids, shard_index, src_addr) -> Task:
        """低级保存从设备地址"""

    @abstractmethod
    def wait(self, task: Task) -> None:
        """阻塞等待任务完成"""

    @abstractmethod
    def check(self, task: Task) -> bool:
        """非阻塞检查任务状态"""

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

# sparse/base.py 核心接口

class UcmSparseRole(Enum):
    SCHEDULER = 0  # 调度器进程
    WORKER = 1     # Worker进程

class UcmSparseMetadata(ABC):
    """Scheduler到Worker的稀疏元数据"""
    pass

class UcmSparseBase(ABC):
    """稀疏注意力基类"""

    def __init__(self, vllm_config, role):
        self._vllm_config = vllm_config
        self._role = role

    # === Worker端方法 ===

    def bind_sparse_metadata(self, metadata: UcmSparseMetadata):
        """绑定从Scheduler传来的元数据"""

    def register_kv_caches(self, kv_caches: Dict[str, torch.Tensor]):
        """注册KV Cache张量"""

    def execute_begin(self, scheduler_output):
        """模型执行开始前"""

    def execute_finished(self, logits_indices) -> torch.Tensor:
        """模型执行完成后"""

    def attention_begin(self, query, key, value, layer_name, ...):
        """注意力计算开始前"""

    def attention_finished(self, query, key, value, attn_output, ...):
        """注意力计算完成后"""

    # === Scheduler端方法 ===

    @abstractmethod
    def request_begin(self, request_id, prompt_token_ids):
        """请求开始"""

    def estimate_num_slots_sparsed(self, request) -> int:
        """估计稀疏后需要的槽位数"""

    def build_sparse_meta(self, scheduler_output, ...) -> UcmSparseMetadata:
        """构建稀疏元数据"""

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

class Scheduler:
    def schedule(self) -> SchedulerOutputs:
        # 1. 处理running请求
        for seq_group in self.running:
            # UCM: 查询外部命中
            if has_kv_connector():
                external_tokens = connector.get_num_new_matched_tokens(
                    seq_group.request,
                    seq_group.num_computed_tokens
                )
                # 更新可跳过的计算量

        # 2. 调度新请求
        scheduled = self._schedule_prefills()

        # 3. 构建输出
        scheduler_output = SchedulerOutput(...)

        # UCM: 附加连接器元数据
        if has_kv_connector():
            scheduler_output.kv_connector_metadata = \
                connector.build_connector_meta(scheduler_output)

        # UCM: 附加稀疏元数据
        if has_ucm_sparse():
            scheduler_output.sparse_metadata = \
                sparse.build_sparse_meta(scheduler_output, ...)

        return scheduler_output

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

class ModelRunner:
    def execute_model(self, scheduler_output):
        # 1. 绑定元数据
        if has_kv_connector():
            connector.bind_connector_metadata(
                scheduler_output.kv_connector_metadata
            )

        if has_ucm_sparse():
            sparse.bind_sparse_metadata(
                scheduler_output.sparse_metadata
            )

        # 2. 开始异步加载
        if has_kv_connector():
            connector.start_load_kv(self.forward_context)

        # UCM: 执行开始钩子
        if has_ucm_sparse():
            sparse.execute_begin(scheduler_output)

        # 3. 模型前向传播
        with set_forward_context(self.attn_metadata):
            hidden_states = self.model.forward(input_ids, positions)

        # UCM: 执行完成钩子
        if has_ucm_sparse():
            logits_indices = sparse.execute_finished(logits_indices)

        # 4. 等待保存完成
        if has_kv_connector():
            connector.wait_for_save()

        return ModelOutput(...)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

# ucm_config.yaml

# 连接器配置
ucm_connectors:
  - ucm_connector_name: "UcmNfsStore"
    ucm_connector_config:
      storage_backends: "/mnt/nfs/kv_cache"
      io_direct: false
      stream_number: 32
      buffer_number: 512
      block_size: 16

# 稀疏注意力配置
ucm_sparse_config:
  ESA:
    init_window_sz: 1
    sparse_ratio: 0.3
    retrieval_stride: 16
  # 或者
  GSA:
    topk_ratio: 0.3
    init_window_size: 2
    local_window_size: 4
  # 或者
  GSAOnDevice:
    hash_bits: 256
    topk_tokens: 1024

# 监控配置
metrics_config_path: "/etc/ucm/metrics.yaml"

# 优化选项
load_only_first_rank: false
enable_prefix_cache: true

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

class UCMConfig:
    """UCM配置管理"""

    @classmethod
    def from_yaml(cls, path: str) -> 'UCMConfig':
        """从YAML文件加载配置"""
        with open(path) as f:
            raw = yaml.safe_load(f)

        return cls(
            connectors=cls._parse_connectors(raw.get('ucm_connectors', [])),
            sparse_config=cls._parse_sparse(raw.get('ucm_sparse_config', {})),
            metrics_path=raw.get('metrics_config_path'),
            load_only_first_rank=raw.get('load_only_first_rank', False),
        )

    @classmethod
    def from_vllm_config(cls, vllm_config) -> 'UCMConfig':
        """从vLLM配置提取UCM配置"""
        extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config

        if 'ucm_config_path' in extra_config:
            return cls.from_yaml(extra_config['ucm_config_path'])
        else:
            return cls._parse_inline(extra_config)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

# 平台选择
export PLATFORM=cuda  # cuda, ascend, musa, maca

# 稀疏模块编译
export ENABLE_SPARSE=true

# 调试选项
export UCM_LOG_LEVEL=DEBUG
export UCM_METRICS_ENABLE=true

# 存储配置覆盖
export UCM_STORAGE_PATH=/mnt/nfs/kv_cache
export UCM_BLOCK_SIZE=32

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

# ucm/logger.py

import logging
from logging.handlers import RotatingFileHandler

def init_logger(name: str = "ucm",
                level: str = "INFO",
                log_file: str = None) -> logging.Logger:
    """初始化UCM日志器"""

    logger = logging.getLogger(name)
    logger.setLevel(getattr(logging, level.upper()))

    # 控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(logging.Formatter(
        '[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s'
    ))
    logger.addHandler(console_handler)

    # 文件处理器（可选）
    if log_file:
        file_handler = RotatingFileHandler(
            log_file, maxBytes=100*1024*1024, backupCount=5
        )
        file_handler.setFormatter(logging.Formatter(
            '[%(asctime)s] [%(levelname)s] [%(name)s] '
            '[%(filename)s:%(lineno)d] %(message)s'
        ))
        logger.addHandler(file_handler)

    return logger

# 全局日志器
ucm_logger = init_logger()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

# shared/metrics/observability.py

from prometheus_client import Counter, Histogram, Gauge

class UCMMetrics:
    """UCM监控指标"""

    # 请求计数
    requests_total = Counter(
        'ucm_requests_total',
        'Total number of requests',
        ['status']  # success, error
    )

    # 缓存命中率
    cache_hit_rate = Gauge(
        'ucm_cache_hit_rate',
        'Cache hit rate',
        ['cache_type']  # prefix, block
    )

    # I/O延迟
    io_latency = Histogram(
        'ucm_io_latency_seconds',
        'I/O operation latency',
        ['operation'],  # load, dump
        buckets=[.001, .005, .01, .05, .1, .5, 1.0]
    )

    # 吞吐量
    io_throughput = Gauge(
        'ucm_io_throughput_gbps',
        'I/O throughput in GB/s',
        ['operation']
    )

    # 内存使用
    memory_usage = Gauge(
        'ucm_memory_usage_bytes',
        'Memory usage',
        ['location']  # gpu, cpu, external
    )

    @classmethod
    def record_load(cls, duration: float, size: int):
        """记录加载操作"""
        cls.io_latency.labels(operation='load').observe(duration)
        cls.io_throughput.labels(operation='load').set(
            size / duration / 1e9  # GB/s
        )

    @classmethod
    def record_hit_rate(cls, hits: int, total: int, cache_type: str):
        """记录命中率"""
        rate = hits / total if total > 0 else 0
        cls.cache_hit_rate.labels(cache_type=cache_type).set(rate)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

# metrics/metrics_configs.yaml

prometheus:
  scrape_interval: 15s
  metrics_path: /metrics
  port: 9090

grafana:
  dashboards:
    - name: UCM Overview
      panels:
        - title: Cache Hit Rate
          type: gauge
          query: ucm_cache_hit_rate

        - title: I/O Throughput
          type: timeseries
          query: ucm_io_throughput_gbps

        - title: Request Latency (p99)
          type: timeseries
          query: histogram_quantile(0.99, ucm_io_latency_seconds_bucket)

        - title: Memory Distribution
          type: piechart
          query: ucm_memory_usage_bytes

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

@abstractmethod
def lookup(self, block_ids: List[bytes]) -> List[bool]:
    """
    查询块是否存在于存储中

    设计要点:
    1. 批量操作：一次查询多个块
    2. 快速返回：只检查元数据，不读取数据
    3. 顺序保持：返回结果与输入顺序一致

    使用场景:
    - Scheduler查询外部命中
    - 判断是否需要计算

    Args:
        block_ids: MD5哈希的块ID列表，每个16字节

    Returns:
        布尔列表，True表示块存在
    """
    pass

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

@abstractmethod
def load_data(
    self,
    block_ids: List[bytes],
    shard_index: List[int],
    dst_addr: List[List[int]] | np.ndarray,
) -> Task:
    """
    从存储加载数据到设备地址

    设计要点:
    1. 直接地址：使用设备指针，避免中间拷贝
    2. 分片支持：shard_index用于张量并行
    3. 异步操作：返回Task句柄

    内存布局:
    dst_addr[i][j] = 第i个块在第j层的目标地址

    Args:
        block_ids: 块哈希列表
        shard_index: 每个块的分片索引
        dst_addr: 目标设备地址矩阵

    Returns:
        异步任务句柄
    """
    pass

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

class Task(ABC):
    """
    异步任务抽象

    生命周期:
    1. 创建：由load_data/dump_data返回
    2. 进行中：I/O在后台执行
    3. 完成：check()返回True或wait()返回

    使用模式:

    # 模式1：阻塞等待
    task = store.load_data(...)
    store.wait(task)  # 阻塞直到完成

    # 模式2：轮询
    task = store.load_data(...)
    while not store.check(task):
        do_other_work()

    # 模式3：批量等待
    tasks = [store.load_data(...) for _ in range(N)]
    for task in tasks:
        store.wait(task)
    """
    pass

# 具体实现示例
class TaskImpl:
    def __init__(self, task_id: int, store_ref):
        self.task_id = task_id
        self.store_ref = store_ref
        self._completed = False

    @property
    def completed(self) -> bool:
        if not self._completed:
            self._completed = self.store_ref.check(self.task_id)
        return self._completed

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

# store/ucmstore.py - V0接口（保留兼容）

class UcmKVStoreBase(ABC):
    """V0存储接口 - 简化版"""

    @abstractmethod
    def create(self, block_ids: List[bytes]) -> List[int]:
        """在存储中预留块空间"""

    @abstractmethod
    def lookup(self, block_ids: List[bytes]) -> List[bool]:
        """查询块存在"""

    @abstractmethod
    def prefetch(self, block_ids: List[bytes]):
        """预取"""

    @abstractmethod
    def load(self, block_ids, offset, dst_tensor) -> Task:
        """加载到张量"""

    @abstractmethod
    def dump(self, block_ids, offset, src_tensor) -> Task:
        """保存从张量"""

    @abstractmethod
    def commit(self, block_ids, is_success: bool):
        """提交或回滚"""

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

# V0到V1的适配器
class V0ToV1Adapter(UcmKVStoreBaseV1):
    """将V0接口适配为V1接口"""

    def __init__(self, v0_store: UcmKVStoreBase):
        self.v0_store = v0_store

    def cc_store(self) -> int:
        return 0  # V0不支持

    def lookup(self, block_ids):
        return self.v0_store.lookup(block_ids)

    def load_data(self, block_ids, shard_index, dst_addr):
        # 将地址转换为张量
        tensors = self._addrs_to_tensors(dst_addr)
        return self.v0_store.load(block_ids, [0]*len(block_ids), tensors)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

# store/factory_v1.py

class UcmConnectorFactoryV1:
    """V1存储连接器工厂"""

    _registry: Dict[str, Tuple[str, str]] = {}  # name -> (module_path, class_name)

    @classmethod
    def register_connector(cls, name: str, module_path: str, class_name: str):
        """注册存储连接器"""
        cls._registry[name] = (module_path, class_name)

    @classmethod
    def create_connector(
        cls,
        connector_name: str,
        config: Dict[str, object]
    ) -> UcmKVStoreBaseV1:
        """创建存储连接器实例"""
        if connector_name not in cls._registry:
            raise ValueError(f"Unknown connector: {connector_name}")

        module_path, class_name = cls._registry[connector_name]

        # 延迟导入模块
        module = importlib.import_module(module_path)
        connector_cls = getattr(module, class_name)

        return connector_cls(config)

# 注册默认存储
UcmConnectorFactoryV1.register_connector(
    "UcmNfsStore",
    "ucm.store.nfsstore.nfsstore_connector",
    "UcmNfsStore"
)

UcmConnectorFactoryV1.register_connector(
    "UcmPcStore",
    "ucm.store.pcstore.pcstore_connector_v1",
    "UcmPcStore"
)

UcmConnectorFactoryV1.register_connector(
    "UcmMooncakeStore",
    "ucm.store.mooncakestore.mooncake_connector",
    "UcmMooncakeStore"
)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

def select_storage(config: Dict) -> UcmKVStoreBaseV1:
    """根据配置选择最合适的存储后端"""

    connectors = config.get('ucm_connectors', [])

    if not connectors:
        raise ValueError("No storage connector configured")

    # 支持多个连接器时的选择逻辑
    for conn_config in connectors:
        name = conn_config['ucm_connector_name']
        params = conn_config['ucm_connector_config']

        # 检查存储路径是否可用
        storage_path = params.get('storage_backends', '')
        if _is_path_available(storage_path):
            return UcmConnectorFactoryV1.create_connector(name, params)

    # 回退到第一个可用的
    first = connectors[0]
    return UcmConnectorFactoryV1.create_connector(
        first['ucm_connector_name'],
        first['ucm_connector_config']
    )

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64

// store/detail/task/task_manager.h

class TaskManager {
public:
    // 提交任务
    Status Submit(Task&& task, size_t& taskId) {
        std::lock_guard<std::mutex> lock(mutex_);

        size_t id = nextTaskId_++;
        tasks_[id] = std::move(task);
        taskId = id;

        // 分配到执行队列
        queues_[id % queues_.size()]->Push(id);

        return Status::OK();
    }

    // 检查任务状态
    Status Check(size_t taskId, bool& finished) {
        std::lock_guard<std::mutex> lock(mutex_);

        auto it = tasks_.find(taskId);
        if (it == tasks_.end()) {
            finished = true;  // 已完成并被回收
            return Status::OK();
        }

        finished = it->second.IsCompleted();
        return Status::OK();
    }

    // 等待任务完成
    Status Wait(size_t taskId) {
        auto startTime = std::chrono::steady_clock::now();

        while (true) {
            bool finished;
            Check(taskId, finished);

            if (finished) {
                // 清理任务
                std::lock_guard<std::mutex> lock(mutex_);
                tasks_.erase(taskId);
                return Status::OK();
            }

            // 超时检查
            auto elapsed = std::chrono::steady_clock::now() - startTime;
            if (elapsed > std::chrono::milliseconds(timeoutMs_)) {
                return Status::Timeout();
            }

            std::this_thread::yield();
        }
    }

private:
    std::unordered_map<size_t, Task> tasks_;
    std::vector<std::unique_ptr<TaskQueue>> queues_;
    std::atomic<size_t> nextTaskId_{0};
    std::mutex mutex_;
    size_t timeoutMs_;
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

import hashlib

def compute_block_hash(parent_hash: bytes, token_ids: List[int]) -> bytes:
    """
    计算块哈希

    选择MD5的理由:
    1. 16字节输出足够用作块ID
    2. 计算速度快
    3. Python标准库支持
    4. 碰撞概率对此场景可接受
    """
    content = parent_hash + bytes(token_ids)
    return hashlib.md5(content).digest()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

class BlockId:
    """块ID的存储和操作"""

    SIZE = 16  # MD5输出16字节

    def __init__(self, hash_bytes: bytes):
        assert len(hash_bytes) == self.SIZE
        self.bytes = hash_bytes

    def to_hex(self) -> str:
        """转换为32字符十六进制字符串"""
        return self.bytes.hex()

    @classmethod
    def from_hex(cls, hex_str: str) -> 'BlockId':
        """从十六进制字符串创建"""
        return cls(bytes.fromhex(hex_str))

    def __hash__(self) -> int:
        """支持作为字典键"""
        return hash(self.bytes)

    def __eq__(self, other) -> bool:
        return self.bytes == other.bytes

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

# store/posix/connector.py

class UcmPosixStore(UcmKVStoreBaseV1):
    """
    基于POSIX的本地文件系统存储

    特点:
    - 零依赖：只需要本地文件系统
    - 低延迟：直接I/O，无网络开销
    - 简单可靠：适合单机场景

    文件布局:
    storage_path/
    ├── blocks/
    │   ├── 00/  # 按哈希前2字符分桶
    │   │   ├── 00a1b2c3...ef.bin
    │   │   └── 00d4e5f6...ab.bin
    │   ├── 01/
    │   └── ...
    └── meta/
        └── index.db  # 块索引（可选）
    """

    def __init__(self, config: Dict):
        self.storage_path = config['storage_backends']
        self.block_size = config.get('block_size', 16)
        self.io_direct = config.get('io_direct', False)
        self.device_id = config.get('device', 0)

        # 初始化C++存储
        self.store = self._init_cpp_store()

    def _init_cpp_store(self):
        """初始化C++底层存储"""
        from ucm.store.posix.cc import PosixStore

        cfg = PosixStoreConfig()
        cfg.storagePath = self.storage_path
        cfg.blockSize = self.block_size
        cfg.ioDirect = self.io_direct
        cfg.deviceId = self.device_id

        store = PosixStore()
        store.Setup(cfg)
        return store

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

// store/posix/cc/posix_file.cpp

class PosixFile {
public:
    Status OpenDirect(const std::string& path, int flags) {
        // 使用O_DIRECT绕过页缓存
        int fd = open(path.c_str(), flags | O_DIRECT);
        if (fd < 0) {
            return Status::IOError(strerror(errno));
        }
        fd_ = fd;
        return Status::OK();
    }

    Status ReadAligned(void* buffer, size_t size, off64_t offset) {
        // O_DIRECT要求对齐
        // 缓冲区地址必须512字节对齐
        // 读取大小必须是512的倍数
        assert(reinterpret_cast<uintptr_t>(buffer) % 512 == 0);
        assert(size % 512 == 0);
        assert(offset % 512 == 0);

        ssize_t ret = pread(fd_, buffer, size, offset);
        if (ret != static_cast<ssize_t>(size)) {
            return Status::IOError("Short read");
        }
        return Status::OK();
    }

private:
    int fd_ = -1;
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

// store/nfsstore/cc/domain/space/space_manager.h

class SpaceManager {
public:
    // 块分配
    std::vector<BlockInfo> Allocate(const std::vector<BlockId>& blockIds) {
        std::vector<BlockInfo> result;

        for (const auto& id : blockIds) {
            // 1. 检查是否已存在
            if (auto info = Lookup(id)) {
                result.push_back(*info);
                continue;
            }

            // 2. 分配新空间
            BlockInfo info;
            info.id = id;
            info.path = ComputePath(id);
            info.offset = 0;
            info.size = blockDataSize_;

            // 3. 确保目录存在
            EnsureDirectory(info.path);

            // 4. 记录元数据
            Register(info);
            result.push_back(info);
        }

        return result;
    }

    // 块查找
    std::vector<bool> Lookup(const std::vector<BlockId>& blockIds) {
        std::vector<bool> exists;

        for (const auto& id : blockIds) {
            // 检查本地索引
            if (index_.contains(id)) {
                exists.push_back(true);
            } else {
                // 检查文件系统
                std::string path = ComputePath(id);
                exists.push_back(FileExists(path));
            }
        }

        return exists;
    }

private:
    std::string ComputePath(const BlockId& id) {
        // 按哈希前缀分桶
        std::string hex = id.ToHex();
        return basePath_ + "/" + hex.substr(0, 2) + "/" + hex + ".bin";
    }

    std::unordered_map<BlockId, BlockInfo> index_;
    std::string basePath_;
    size_t blockDataSize_;
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

# NFS Store支持多个存储后端

config = {
    "storage_backends": "/mnt/nfs1:/mnt/nfs2:/mnt/ssd_cache",
    # 冒号分隔的多个路径

    "storage_policy": "tiered",  # tiered | round_robin | random
}

# 分层策略:
# /mnt/ssd_cache: 热数据，最快访问
# /mnt/nfs1:      温数据，中速访问
# /mnt/nfs2:      冷数据，备份存储

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

# store/cache/connector.py

class UcmCacheStore(UcmKVStoreBaseV1):
    """
    高速缓存层

    架构:
    ┌─────────────────────────────┐
    │    GPU/CPU 高速缓存          │
    │    (CacheStore)              │
    └──────────────┬──────────────┘
                   │ 未命中时
                   ▼
    ┌─────────────────────────────┐
    │    后端存储                  │
    │    (POSIX/NFS/DS3FS)        │
    └─────────────────────────────┘
    """

    def __init__(self, config: Dict):
        # 后端存储
        backend_config = config['backend']
        self.backend = UcmConnectorFactoryV1.create_connector(
            backend_config['name'],
            backend_config['config']
        )

        # 缓存配置
        self.cache_size = config.get('cache_size', 1 << 30)  # 1GB
        self.device_id = config.get('device', 0)

        # C++缓存实现
        self.cache = self._init_cache()

    def lookup(self, block_ids: List[bytes]) -> List[bool]:
        # 先查缓存
        cache_results = self.cache.Lookup(block_ids)

        # 缓存未命中的查后端
        backend_queries = [
            bid for bid, hit in zip(block_ids, cache_results) if not hit
        ]

        if backend_queries:
            backend_results = self.backend.lookup(backend_queries)
            # 合并结果
            ...

        return results

    def load_data(self, block_ids, shard_index, dst_addr):
        # 分离缓存命中和未命中
        cache_hits = []
        backend_loads = []

        for i, bid in enumerate(block_ids):
            if self.cache.Contains(bid):
                cache_hits.append((i, bid))
            else:
                backend_loads.append((i, bid))

        # 从缓存加载
        if cache_hits:
            cache_task = self.cache.Load(
                [bid for _, bid in cache_hits],
                [dst_addr[i] for i, _ in cache_hits]
            )

        # 从后端加载并填充缓存
        if backend_loads:
            # 加载到临时缓冲
            temp_addrs = self._alloc_temp_buffer(len(backend_loads))
            backend_task = self.backend.load_data(
                [bid for _, bid in backend_loads],
                [shard_index[i] for i, _ in backend_loads],
                temp_addrs
            )

            # 等待并填充缓存
            self.backend.wait(backend_task)
            self.cache.Insert(
                [bid for _, bid in backend_loads],
                temp_addrs
            )

            # 复制到目标
            self._copy_to_dst(temp_addrs,
                             [dst_addr[i] for i, _ in backend_loads])

        return CompositeTask([cache_task, backend_task])

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

// store/cache/cc/eviction_policy.h

class LRUEvictionPolicy {
public:
    void Touch(const BlockId& id) {
        std::lock_guard<std::mutex> lock(mutex_);

        // 移动到队尾（最近使用）
        auto it = position_.find(id);
        if (it != position_.end()) {
            lru_queue_.splice(lru_queue_.end(), lru_queue_, it->second);
        }
    }

    std::vector<BlockId> Evict(size_t count) {
        std::lock_guard<std::mutex> lock(mutex_);

        std::vector<BlockId> evicted;
        for (size_t i = 0; i < count && !lru_queue_.empty(); ++i) {
            BlockId id = lru_queue_.front();
            lru_queue_.pop_front();
            position_.erase(id);
            evicted.push_back(id);
        }

        return evicted;
    }

private:
    std::list<BlockId> lru_queue_;
    std::unordered_map<BlockId, std::list<BlockId>::iterator> position_;
    std::mutex mutex_;
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

# store/ds3fs/connector.py

class UcmDs3fsStore(UcmKVStoreBaseV1):
    """
    DS3FS分布式存储

    特点:
    - 高吞吐: 聚合带宽可达TB/s级别
    - 低延迟: 优化的网络栈
    - 高可用: 多副本冗余

    配置:
    - ior_entries: I/O请求条目数
    - ior_depth: I/O请求深度
    - numa_id: NUMA节点绑定
    """

    def __init__(self, config: Dict):
        self.mount_point = config['storage_backends']
        self.ior_entries = config.get('ior_entries', 1024)
        self.ior_depth = config.get('ior_depth', 32)
        self.numa_id = config.get('numa_id', 0)

        # 初始化DS3FS客户端
        self.client = self._init_client()