clean code

Author: wuhuxiao

docs/source/user-guide/sparse-attention/cacheblend.md

@@ -29,10 +29,10 @@ CacheBlend reduces TTFT by 2.2 ~ 3.3× and increases throughput by 2.8 ~ 5× und
 
 ### Native Block-Wise Chunk KV Cache Dump, Load, PostProcess and Recompute
 1. **🔐 Chunk Hash Encoding**: Similar as prefix hash encoder, hash all blocks in each chunk from the same hash meta beginning.
-2. **⚡ Combine Prefix Cache and Chunk Cache**: Since chunk cache and native prefix cache share the same hash space, ucm first performs prefix cache lookup to fetch fully resued cache and then conduct chunk cache lookup to fetch the candidate cache for blending.
+2. **⚡ Combine Prefix Cache and Chunk Cache**: Since chunk cache and native prefix cache share the same hash space, ucm first performs prefix cache lookup to fetch fully reused cache and then conduct chunk cache lookup to fetch the candidate cache for blending.
 3. **🎯 Delta-Rope PostProcess**: Rectify loaded chunk cache according to their position in the new request.
-3. **🔍 Integrate Cache Blend and First Token Generation**: Construct compute mask and attention meta according to HKVD tokens, cache miss tokens and suffix tokens, then compute their kv cache in a single model forward stage 
-4. **🚀 Comprehensive Hook for LLM Forward Pipeline**: Based on ucm sparse module, blend module spare the prefill tokens not only in attention stage but also in ffn, layer stage.
+3. **🔍 Integrate Cache Blend and First Token Generation**: Construct compute mask and attention meta according to HKVD tokens, cache miss tokens and suffix tokens, then compute their kv cache in a single model forward stage.
+4. **🚀 Comprehensive Hook for LLM Forward Pipeline**: Based on ucm sparse module, blend module sparse the prefill tokens not only in attention stage but also in ffn, layer stage.
 
 ## 🚀 Quick Start
 
@@ -49,7 +49,7 @@ python <ucm-repo>/examples/offline_inference_blend.py
 ```
 
 ### Basic Usage
-Similr to UCM's `offline_inference_esa.py` examples. We only need to specify `ucm_sparse_method` to be `Blend` and specify meta config, as shown below.
+Similar to UCM's `offline_inference_esa.py` examples. We only need to specify `ucm_sparse_method` to be `Blend` and specify meta config, as shown below.
 
 ```python
 ...

ucm/integration/vllm/blend_connector.py

@@ -32,10 +32,10 @@
 @dataclass
 class ChunkMetaData:
     # [start, start + len)
-    start_idx_in_req: int
+    start_token_dix: int
     chunk_tokens_len: int
 
-    start_idx_in_req_blks: int
+    start_blk_idx: int
     chunk_blks_len: int
 
     cached_start_position: int
@@ -45,20 +45,20 @@ class ChunkMetaData:
     store_hits: List[bool] = field(default_factory=list)
 
     @property
-    def end_idx_in_req(self) -> int:
-        return self.start_idx_in_req + self.chunk_tokens_len
+    def end_token_dix(self) -> int:
+        return self.start_token_dix + self.chunk_tokens_len
 
     @property
-    def end_idx_in_req_blks(self) -> int:
-        return self.start_idx_in_req_blks + self.chunk_blks_len
+    def end_blk_idx(self) -> int:
+        return self.start_blk_idx + self.chunk_blks_len
 
     @property
     def cached_end_position(self) -> int:
         return self.cached_start_position + self.chunk_tokens_len
 
     @property
     def position_offset(self) -> int:
-        return self.start_idx_in_req - self.cached_start_position
+        return self.start_token_dix - self.cached_start_position
 
     @property
     def hits_vllm_blk_ids(self) -> List[int]:
@@ -77,10 +77,10 @@ def merge_chunk(self, temp_chunk_meta: Self) -> None:
 
     def update_meta_partial_pc(self, num_pc_part_blks: int, block_size: int) -> None:
         if num_pc_part_blks > 0:
-            self.start_idx_in_req += num_pc_part_blks * block_size
+            self.start_token_dix += num_pc_part_blks * block_size
             self.chunk_tokens_len -= num_pc_part_blks * block_size
 
-            self.start_idx_in_req_blks += num_pc_part_blks
+            self.start_blk_idx += num_pc_part_blks
             self.chunk_blks_len -= num_pc_part_blks
 
             self.chunk_blks_hash = self.chunk_blks_hash[num_pc_part_blks:]
@@ -211,9 +211,9 @@ def _process_req(self, all_token_ids: List[int]):
                 chunk_tokens_len = chunk_blks_len * self.block_size
 
                 rag_chunk_meta = ChunkMetaData(
-                    start_idx_in_req=start_token_dix,
+                    start_token_dix=start_token_dix,
                     chunk_tokens_len=chunk_tokens_len,
-                    start_idx_in_req_blks=start_blk_idx,
+                    start_blk_idx=start_blk_idx,
                     chunk_blks_len=chunk_blks_len,
                     chunk_blks_hash=chunk_blks_hash,
                     cached_start_position=0,
@@ -271,7 +271,7 @@ def _get_req_chunk_hit(
         # for cache blend
         for i, chunk_meta in enumerate(req_chunks_meta):
             chunk_meta.store_hits = chunk_lookup_results[
-                chunk_meta.start_idx_in_req_blks : chunk_meta.end_idx_in_req_blks
+                chunk_meta.start_blk_idx : chunk_meta.end_blk_idx
             ]
         first_chunk_meta = req_chunks_meta[0]
         first_chunk_meta.update_meta_partial_pc(pc_hit_blocks, self.block_size)
@@ -324,7 +324,7 @@ def _generate_blend_dispatch_meta(
             # just need to load, in future we may create a multi-chunk hash to dump and reuse the blended cache
             for chunk_meta in req_meta.chunks_meta:
                 chunk_meta.vllm_blk_ids = vllm_block_ids[
-                    chunk_meta.start_idx_in_req_blks : chunk_meta.end_idx_in_req_blks
+                    chunk_meta.start_blk_idx : chunk_meta.end_blk_idx
                 ]
                 load_ucm_block_ids.extend(chunk_meta.hits_chunk_blks_hash)
                 load_vllm_block_ids.extend(chunk_meta.hits_vllm_blk_ids)

ucm/sparse/blend/blend.py

@@ -85,10 +85,8 @@ def add_request(
                 hit_mask.extend(meta.store_hits)
             reqMeta = ReqMeta(
                 req_idx=req_idx_batch,
-                prefix_len=chunks_meta[0].start_idx_in_req,
-                prefix_blk_len=get_num_blks(
-                    chunks_meta[0].start_idx_in_req, block_size
-                ),
+                prefix_len=chunks_meta[0].start_token_dix,
+                prefix_blk_len=get_num_blks(chunks_meta[0].start_token_dix, block_size),
                 chunks_len=len(hit_mask) * block_size,
                 chunks_blk_len=len(hit_mask),
                 chunk_hit_mask=hit_mask,