clean code

Author: wuhuxiao

ucm/integration/vllm/blend_connector.py

@@ -1,21 +1,14 @@
-import hashlib
 import itertools
-import os
-import pickle
-import time
 from dataclasses import dataclass, field
 from enum import Enum, auto
-from typing import TYPE_CHECKING, Callable, List, Optional, Self, Tuple
+from typing import TYPE_CHECKING, List, Self, Tuple
 
 import torch
 from vllm.config import VllmConfig
 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
-    KVConnectorBase_V1,
     KVConnectorMetadata,
     KVConnectorRole,
 )
-from vllm.distributed.parallel_state import get_tp_group, get_world_group
-from vllm.platforms import current_platform
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.request import Request
 
@@ -28,16 +21,9 @@
 )
 from ucm.logger import init_logger
 from ucm.shared.metrics import ucmmonitor
-from ucm.shared.metrics.observability import UCMStatsLogger
 from ucm.sparse.blend.blockwise_rope import block_wise_rope_forward
-from ucm.sparse.kvstar.multistep import ReqStage
-from ucm.store.factory import UcmConnectorFactory
-from ucm.store.ucmstore import Task, UcmKVStoreBase
-from ucm.utils import Config
 
 if TYPE_CHECKING:
-    from vllm.attention.backends.abstract import AttentionMetadata
-    from vllm.forward_context import ForwardContext
     from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 
 logger = init_logger(__name__)
@@ -82,7 +68,7 @@ def hits_vllm_blk_ids(self) -> List[int]:
     def hits_chunk_blks_hash(self) -> List[str]:
         return list(itertools.compress(self.chunk_blks_hash, self.store_hits))
 
-    def merge_chunk(self, temp_chunk_meta: Self):
+    def merge_chunk(self, temp_chunk_meta: Self) -> None:
         # current we use a fix pattern(end with a fix token id) to recognize the text token chunk
         # in some special situation, one text chunk maybe split as multi text chunk, so we should merge them into one
         self.chunk_tokens_len += temp_chunk_meta.chunk_tokens_len
@@ -107,10 +93,10 @@ class BlendStage(Enum):
     BUILD_PREFIX_CACHE = auto()
     CACHE_BLEND = auto()
 
-    def is_blend_cache(self):
+    def is_blend_cache(self) -> bool:
         return self == BlendStage.CACHE_BLEND
 
-    def is_prefix_cache(self):
+    def is_prefix_cache(self) -> bool:
         return self == BlendStage.BUILD_PREFIX_CACHE
 
 
@@ -137,10 +123,7 @@ class UCMBlendConnectorMetadata(UCMConnectorMetadata):
 
 class UCMBlendConnector(UCMDirectConnector):
     """
-    This Connector means overlap:
-    load l0 -> forward l0 -> save l0
-               load l1    -> forward l1 -> save l1
-                             load l2    -> forward l2 -> save l2
+    This Connector process chunk hash and prefix cache
     """
 
     def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
@@ -265,7 +248,7 @@ def _get_req_chunk_hit(
         prefix_block_hashes: List[str],
         req_chunks_meta: List[ChunkMetaData],
         req_chunks_hashes: List[str],
-    ):
+    ) -> Tuple[int, int]:
 
         # first perform prefix cache lookup
         pc_lookup_results = self.store.lookup(prefix_block_hashes)
@@ -312,7 +295,7 @@ def _generate_blend_dispatch_meta(
         ----------------------------------------------------------------------------------------------------------
         |            LOAD               |          DUMP          |
         ----------------------------------------------------------------------------------------------------------
-        |           REUSE               |    RECOMPUTE           |
+        |           REUSE               |     RECOMPUTE          |
         ----------------------------------------------------------------------------------------------------------
 
 
@@ -362,7 +345,7 @@ def _generate_blend_dispatch_meta(
             req_meta.chunks_meta,
         )
 
-    def _post_process_chunk_cache(self, k_cache, vllm_ids, positions):
+    def _post_process_chunk_cache(self, k_cache, vllm_ids, positions) -> None:
         """
         post process loaded chunk kcache
         """
@@ -371,7 +354,7 @@ def _post_process_chunk_cache(self, k_cache, vllm_ids, positions):
         # triton kernl for block-wise delta rope
         block_wise_rope_forward(k_cache, vllm_ids, positions, self.cos_sin_cache)
 
-    def _register_cos_sin_cache(self, model: "Model"):
+    def _register_cos_sin_cache(self, model: "Model") -> None:
         try:
             rotary_emb = model.model.layers[0].self_attn.rotary_emb
             self.cos_sin_cache = rotary_emb.cos_sin_cache

ucm/sparse/blend/README.md

@@ -13,10 +13,9 @@
 
 from vllm.config import VllmConfig
 from vllm.forward_context import ForwardContext
-from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.request import Request
 
-from ucm.integration.vllm.blend_connector import BlendRequestDispatchMeta, ChunkMetaData
+from ucm.integration.vllm.blend_connector import BlendRequestDispatchMeta
 from ucm.sparse.base import (
     INVALID_SLOT,
     UcmSparseBase,

ucm/sparse/blend/blend.py

@@ -59,7 +59,12 @@ def _triton_rope_blockwise_kernel(
     tl.store(k_ptr + offs + hd // 2, new_k_tile_2, mask=mask)
 
 
-def block_wise_rope_forward(k_cache, vllm_ids, positions, cos_sin_cache):
+def block_wise_rope_forward(
+    k_cache: torch.Tensor,
+    vllm_ids: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+) -> torch.Tensor:
     """
     Args：
         k_cache: torch.Tensor (total_blocks, seq_len, n_kv_heads, hd), vllm owned.
@@ -96,7 +101,12 @@ def block_wise_rope_forward(k_cache, vllm_ids, positions, cos_sin_cache):
     return k_cache
 
 
-def rope_naive_torch(k_cache, vllm_ids, positions, cos_sin_cache):
+def rope_naive_torch(
+    k_cache: torch.Tensor,
+    vllm_ids: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+) -> torch.Tensor:
     """
     naive torch implementation for accuracy and perf baseline
     Args:

ucm/sparse/blend/blockwise_rope.py

@@ -933,6 +933,7 @@ def execute_finished(self, logits_indices: torch.Tensor):
             self.prefetch_engine.deal_async_prefetch(
                 False, self.gsa_metadata, kv_caches, None
             )
+        return logits_indices
 
     def launch_transfer_task(self, all_free_block_ids, all_miss_ids, kv_caches):
         if all_free_block_ids == None:
@@ -1006,8 +1007,6 @@ def check_transfer_task_done(self) -> bool:
         self.task_load.clear()
         return True
 
-        return logits_indices
-
     def build_sparse_meta(
         self, scheduler_output: SchedulerOutput, requests, input_batch, attn_metadata
     ) -> None:

ucm/sparse/blend/utils.py

@@ -82,7 +82,6 @@ def maybe_execute_sparse_layer_begin(
     if not has_ucm_sparse():
         return positions, hidden_states, residual
     ucm_spare = get_ucm_sparse()
-    # after sparse, n_tokens of source tensor is larger than target
     return ucm_spare.layer_begin(positions, hidden_states, residual)
 
 
@@ -92,15 +91,13 @@ def maybe_execute_sparse_layer_finished(
     if not has_ucm_sparse():
         return positions, hidden_states, residual
     ucm_spare = get_ucm_sparse()
-    # after sparse, n_tokens of source tensor is larger than target
     return ucm_spare.layer_finished(positions, hidden_states, residual)
 
 
 def maybe_execute_sparse_ffn_begin(hidden_states: torch.Tensor, residual: torch.Tensor):
     if not has_ucm_sparse():
         return hidden_states, residual
     ucm_spare = get_ucm_sparse()
-    # after sparse, n_tokens of source tensor is larger than target
     return ucm_spare.ffn_begin(hidden_states, residual)
 
 
@@ -110,5 +107,4 @@ def maybe_execute_sparse_ffn_finished(
     if not has_ucm_sparse():
         return hidden_states, residual
     ucm_spare = get_ucm_sparse()
-    # after sparse, n_tokens of source tensor is larger than target
     return ucm_spare.ffn_finished(hidden_states, residual)