Star41
137
代码
代码
Issues155
Pull Requests130
流水线
项目讨论
Wiki
分析
Star41
137
已合并
[bugfix]cpu/disk命中根据parent_hash + token_hash查找 #601
ganglv创建于 27 天前
[bugfix]cpu/disk命中根据parent_hash + token_hash查找 #601
已合并
ganglv创建于 27 天前
从ganglv/MindIE-PyMotor_0701conductor合入到Ascend/MindIE-Motormaster
共 31 个文件变更+2669-1104
M.gitignore+3-0
@@ -127,3 +127,6 @@ motor/kv_conductor/curl.sh
127Dockerfile.motor-e2e-kv127Dockerfile.motor-e2e-kv
128motor/kv_conductor/.cargo/128motor/kv_conductor/.cargo/
129motor/kv_conductor/mock/129motor/kv_conductor/mock/
130+ 
131+# gitleaks
132+gitleaks
Mexamples/deployer/log_collect/log_monitor.py+75-21
@@ -9,7 +9,7 @@
9# See the Mulan PSL v2 for more details.9# See the Mulan PSL v2 for more details.
10 10 
11# Copyright Huawei Technologies Co., Ltd. 2026. All rights reserved.11# Copyright Huawei Technologies Co., Ltd. 2026. All rights reserved.
12-from datetime import datetime, timezone12+from datetime import datetime, timedelta, timezone
13import configparser13import configparser
14import logging14import logging
15import logging.handlers15import logging.handlers
@@ -227,9 +227,24 @@ class LogMonitor:
227 log_e(f"shell_get_pod Exception: {e}")227 log_e(f"shell_get_pod Exception: {e}")
228 return None228 return None
229 229 
230- def shell_pull_log(self, pod_name: str, file_path: str, interval: float = 0.2) -> bool:230+ @staticmethod
231+ def _kubectl_since_time(when: datetime | None = None, lookback_seconds: float = 2.0) -> str:
232+ """RFC3339 UTC timestamp for ``kubectl logs --since-time`` (small lookback avoids gaps)."""
233+ ts = (when or datetime.now(timezone.utc)) - timedelta(seconds=lookback_seconds)
234+ return ts.astimezone(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
235+ 
236+ def shell_pull_log(
237+ self,
238+ pod_name: str,
239+ file_path: str,
240+ interval: float = 0.2,
241+ since_time: str | None = None,
242+ ) -> bool:
231 """243 """
232 Execute the kubectl command to obtain the log.244 Execute the kubectl command to obtain the log.
245+ 
246+ :param since_time: If set, pass ``--since-time`` so reconnects only fetch
247+ new lines (avoids re-dumping the full container history).
233 """248 """
234 b_write_flag = False249 b_write_flag = False
235 abs_path = os.path.abspath(os.path.normpath(file_path))250 abs_path = os.path.abspath(os.path.normpath(file_path))
@@ -240,10 +255,21 @@ class LogMonitor:
240 255 
241 process = None256 process = None
242 try:257 try:
258+ kubectl_args = [
259+ self.cmd_kubectl,
260+ 'logs',
261+ '-f',
262+ '-n',
263+ g_name_space,
264+ pod_name,
265+ ]
266+ if since_time:
267+ kubectl_args.extend(['--since-time', since_time])
268+ 
243 # Long-running kubectl logs -f; with-statement is unsuitable269 # Long-running kubectl logs -f; with-statement is unsuitable
244 # because the process is terminated in the finally block.270 # because the process is terminated in the finally block.
245 process = subprocess.Popen( # pylint: disable=consider-using-with271 process = subprocess.Popen( # pylint: disable=consider-using-with
246- [self.cmd_kubectl, 'logs', '-f', '-n', g_name_space, pod_name],272+ kubectl_args,
247 stdout=subprocess.PIPE,273 stdout=subprocess.PIPE,
248 stderr=subprocess.PIPE,274 stderr=subprocess.PIPE,
249 text=True,275 text=True,
@@ -255,6 +281,8 @@ class LogMonitor:
255 if abs_path not in self._logged_save_paths:281 if abs_path not in self._logged_save_paths:
256 self._logged_save_paths.add(abs_path)282 self._logged_save_paths.add(abs_path)
257 log_i(f"{pod_name}: logs save to: {abs_path} (max {size:.1f}MB, keep {g_backup_count} backups)")283 log_i(f"{pod_name}: logs save to: {abs_path} (max {size:.1f}MB, keep {g_backup_count} backups)")
284+ elif since_time:
285+ log_i(f"{pod_name}: reconnecting logs with --since-time={since_time}")
258 286 
259 # Reads the output in real time and writes it to the log file.287 # Reads the output in real time and writes it to the log file.
260 while not self.exit_flag.is_set():288 while not self.exit_flag.is_set():
@@ -276,7 +304,10 @@ class LogMonitor:
276 # Ensure the child process is terminated304 # Ensure the child process is terminated
277 if process and process.poll() is None:305 if process and process.poll() is None:
278 process.terminate()306 process.terminate()
279- log_i(f"{pod_name} :The thread has exited.")307+ for handler in list(logger.handlers):
308+ handler.close()
309+ logger.removeHandler(handler)
310+ log_i(f"{pod_name}: log stream session ended.")
280 return b_write_flag311 return b_write_flag
281 312 
282 def pull_log_and_save(self, pod_name: str, interval: float = 3) -> None:313 def pull_log_and_save(self, pod_name: str, interval: float = 3) -> None:
@@ -288,6 +319,13 @@ class LogMonitor:
288 next_slot = self._pod_log_next_slot.get(pod_name, 0)319 next_slot = self._pod_log_next_slot.get(pod_name, 0)
289 index = max(generation_floor, next_slot)320 index = max(generation_floor, next_slot)
290 node_name = self.shell_get_pod_node(pod_name)321 node_name = self.shell_get_pod_node(pod_name)
322+ # Reuse one rotating log file across kubectl stream reconnects; only
323+ # allocate a new _{n}.log when the collector starts fresh or the pod
324+ # moves to another node. Otherwise each ~4h apiserver disconnect would
325+ # re-dump full history into _1/_2/... (and duplicate .log.1 rotations).
326+ file_path: str | None = None
327+ allocated_log_index: int | None = None
328+ since_time: str | None = None
291 try:329 try:
292 while not self.exit_flag.is_set():330 while not self.exit_flag.is_set():
293 pod_running_state = self.check_pod_is_running(pod_name)331 pod_running_state = self.check_pod_is_running(pod_name)
@@ -301,26 +339,37 @@ class LogMonitor:
301 fetched_node_name = self.shell_get_pod_node(pod_name)339 fetched_node_name = self.shell_get_pod_node(pod_name)
302 if fetched_node_name != node_name:340 if fetched_node_name != node_name:
303 log_i(f"{pod_name}: node_name refresh {node_name!r} -> {fetched_node_name!r}")341 log_i(f"{pod_name}: node_name refresh {node_name!r} -> {fetched_node_name!r}")
304- node_name = fetched_node_name342+ node_name = fetched_node_name
305- file_path, allocated_log_index = self._allocate_unique_log_path(pod_name, node_name, index)343+ # New node => new file name; fetch full history for this placement.
306- if allocated_log_index != index:344+ if allocated_log_index is not None:
307- log_i(345+ index = max(index, allocated_log_index + 1)
308- f"{pod_name}: log file slot bumped {index} -> {allocated_log_index} "346+ file_path = None
309- "(target path already exists)."347+ allocated_log_index = None
310- )348+ since_time = None
311- if self.shell_pull_log(pod_name, file_path):349+ if file_path is None:
312- next_after = allocated_log_index + 1350+ file_path, allocated_log_index = self._allocate_unique_log_path(pod_name, node_name, index)
351+ if allocated_log_index != index:
352+ log_i(
353+ f"{pod_name}: log file slot bumped {index} -> {allocated_log_index} "
354+ "(target path already exists)."
355+ )
356+ index = allocated_log_index
357+ if self.shell_pull_log(pod_name, file_path, since_time=since_time):
358+ # Reserve next slot for a future collector generation, but keep
359+ # appending to the same file for this thread's reconnects.
313 self._pod_log_next_slot[pod_name] = max(360 self._pod_log_next_slot[pod_name] = max(
314 self._pod_log_next_slot.get(pod_name, 0),361 self._pod_log_next_slot.get(pod_name, 0),
315- next_after,362+ allocated_log_index + 1,
316 )363 )
317- index = self._pod_log_next_slot[pod_name]364+ # Capture before sleep so logs during the pause are still fetched.
318- # Pod restart — apply exponential backoff to avoid log365+ since_time = self._kubectl_since_time()
319- # duplication from frequent restarts (e.g. crash-loop).366+ # Stream reconnect (not a pod crash-loop) — keep a short pause.
320- delay = self._backoff_sleep(pod_name)367+ # Backoff is reserved for cases that start a brand-new file.
368+ delay = interval
321 log_i(369 log_i(
322 f"{pod_name}: Log stream ended; pausing {delay:.0f}s before "370 f"{pod_name}: Log stream ended; pausing {delay:.0f}s before "
323- "re-checking pod and reopening logs if still Running."371+ f"re-checking pod and reopening logs "
372+ f"(same file, --since-time={since_time})."
324 )373 )
325 time.sleep(delay)374 time.sleep(delay)
326 else:375 else:
@@ -329,8 +378,13 @@ class LogMonitor:
329 os.remove(file_path)378 os.remove(file_path)
330 except OSError:379 except OSError:
331 pass380 pass
332- log_w(f"{pod_name}: Failed to pull logs; pausing {interval}s before retry.")381+ file_path = None
333- time.sleep(interval)382+ allocated_log_index = None
383+ since_time = None
384+ # Failed / crash-loop pulls — exponential backoff before retry.
385+ delay = self._backoff_sleep(pod_name)
386+ log_w(f"{pod_name}: Failed to pull logs; pausing {delay:.0f}s before retry.")
387+ time.sleep(delay)
334 except Exception as e:388 except Exception as e:
335 log_e(f"{pod_name} :Exception: {e}")389 log_e(f"{pod_name} :Exception: {e}")
336 finally:390 finally:
Mexamples/features/config_sample.json+1-1
@@ -264,7 +264,7 @@
264 "block_size": 128,264 "block_size": 128,
265 "engine_type": "vLLM",265 "engine_type": "vLLM",
266 "pool_endpoint": "",266 "pool_endpoint": "",
267- "xpu_endpoint": "",267+ "npu_endpoint": "",
268 "cpu_endpoint": "",268 "cpu_endpoint": "",
269 "disk_endpoint": "",269 "disk_endpoint": "",
270 "endpoint": "",270 "endpoint": "",
Mmotor/config/config_utils.py+13-4
@@ -280,21 +280,30 @@ def _update_prefill_kv_event_config(updated_config: dict[str, Any], user_config_
280 280 
281 281 
282def _redirect_prefill_kv_event_config(updated_config: dict[str, Any], user_config_data: dict[str, Any]) -> None:282def _redirect_prefill_kv_event_config(updated_config: dict[str, Any], user_config_data: dict[str, Any]) -> None:
283- """Redirect legacy prefill_kv_event_config into the unified kv_conductor_config.283+ """Redirect legacy / top-level kv conductor settings into scheduler_config.
284 284 
285- 1. If the user config still has a ``prefill_kv_event_config`` section, merge its285+ 1. Merge top-level ``user_config.kv_conductor_config`` into
286+ ``scheduler_config.kv_conductor_config`` (deployer convention).
287+ 2. If the user config still has a ``prefill_kv_event_config`` section, merge its
286 fields into ``kv_conductor_config`` (backward compat).288 fields into ``kv_conductor_config`` (backward compat).
287- 2. Auto-derive connection info from engine sections and kv_conductor_config.289+ 3. Auto-derive connection info from engine sections when fields are unset.
288 """290 """
289 try:291 try:
290 reg = updated_config.setdefault("scheduler_config", {}).setdefault(KV_CONDUCTOR_CONFIG, {})292 reg = updated_config.setdefault("scheduler_config", {}).setdefault(KV_CONDUCTOR_CONFIG, {})
291 293 
294+ # ── Top-level user_config.kv_conductor_config (deployer) ──────
295+ top_level = user_config_data.get(KV_CONDUCTOR_CONFIG)
296+ if isinstance(top_level, dict):
297+ for key, value in top_level.items():
298+ if value not in (None, "") and not reg.get(key):
299+ reg[key] = value
300+ 
292 # ── Backward compat: migrate old prefill_kv_event_config ──────301 # ── Backward compat: migrate old prefill_kv_event_config ──────
293 old_config = updated_config.pop(PREFILL_KV_EVENT_CONFIG, None)302 old_config = updated_config.pop(PREFILL_KV_EVENT_CONFIG, None)
294 if isinstance(old_config, dict):303 if isinstance(old_config, dict):
295 logger.warning(304 logger.warning(
296 "prefill_kv_event_config is deprecated and will be removed in a future version. "305 "prefill_kv_event_config is deprecated and will be removed in a future version. "
297- "Please migrate to kv_conductor_config under scheduler_config. "306+ "Please migrate to top-level kv_conductor_config (or scheduler_config.kv_conductor_config). "
298 "See docs/zh/user_guide/features/kvcache_affinity.md for details."307 "See docs/zh/user_guide/features/kvcache_affinity.md for details."
299 )308 )
300 for key in (309 for key in (
Mmotor/config/coordinator.py+6-3
@@ -157,8 +157,8 @@ class KvConductorConfig:
157 the kv-conductor. Behaviour varies by ``store_backend``:157 the kv-conductor. Behaviour varies by ``store_backend``:
158 158 
159 - Mooncake / Memcache: register the pool once (``pool_endpoint``) +159 - Mooncake / Memcache: register the pool once (``pool_endpoint``) +
160- per-DP HBM via ``xpu_endpoint``.160+ per-DP HBM via ``npu_endpoint``.
161- - YuanRong: per-DP multi-port via ``xpu/cpu/disk_endpoint`` patterns.161+ - YuanRong: per-DP multi-port via ``npu/cpu/disk_endpoint`` patterns.
162 162 
163 Endpoint patterns use ``*`` as IP placeholder and add ``dp_rank``163 Endpoint patterns use ``*`` as IP placeholder and add ``dp_rank``
164 to the port, e.g. ``"tcp://*:15557"`` resolves to164 to the port, e.g. ``"tcp://*:15557"`` resolves to
@@ -195,9 +195,12 @@ class KvConductorConfig:
195 pool_endpoint: str = ""195 pool_endpoint: str = ""
196 """Pool service endpoint for centralized backends, e.g. "tcp://kvp-master:5557"."""196 """Pool service endpoint for centralized backends, e.g. "tcp://kvp-master:5557"."""
197 197 
198- xpu_endpoint: str = ""198+ npu_endpoint: str = ""
199 """Per-DP HBM ZMQ PUB endpoint pattern, e.g. "tcp://*:50090"."""199 """Per-DP HBM ZMQ PUB endpoint pattern, e.g. "tcp://*:50090"."""
200 200 
201+ xpu_endpoint: str = ""
202+ """Deprecated alias of ``npu_endpoint``; used when ``npu_endpoint`` is empty."""
203+ 
201 cpu_endpoint: str = ""204 cpu_endpoint: str = ""
202 """Per-DP CPU/DDR ZMQ PUB endpoint pattern, e.g. "tcp://*:15558"."""205 """Per-DP CPU/DDR ZMQ PUB endpoint pattern, e.g. "tcp://*:15558"."""
203 206 
Mmotor/coordinator/api_client/conductor_api_client.py+12-71
@@ -129,7 +129,11 @@ class ConductorApiClient:
129 def _register_hbm_dp(cls, reg, store_backend: str, instance: "Instance", endpoint: "Endpoint") -> None:129 def _register_hbm_dp(cls, reg, store_backend: str, instance: "Instance", endpoint: "Endpoint") -> None:
130 """Register a single DP's HBM endpoint for pool-backend auto-attach."""130 """Register a single DP's HBM endpoint for pool-backend auto-attach."""
131 instance_id = conductor_instance_id(instance)131 instance_id = conductor_instance_id(instance)
132- xpu_url = cls._resolve_endpoint_url(reg.xpu_endpoint or reg.endpoint, endpoint.ip, endpoint.id)132+ npu_url = cls._resolve_endpoint_url(
133+ reg.npu_endpoint or reg.xpu_endpoint or reg.endpoint,
134+ endpoint.ip,
135+ endpoint.id,
136+ )
133 137 
134 replay_url = cls._resolve_endpoint_url(reg.replay_endpoint, endpoint.ip, endpoint.id)138 replay_url = cls._resolve_endpoint_url(reg.replay_endpoint, endpoint.ip, endpoint.id)
135 register_data: dict = {139 register_data: dict = {
@@ -140,8 +144,8 @@ class ConductorApiClient:
140 "block_size": reg.block_size,144 "block_size": reg.block_size,
141 "dp_rank": endpoint.id,145 "dp_rank": endpoint.id,
142 }146 }
143- if xpu_url:147+ if npu_url:
144- register_data["medium_endpoints"] = {"xpu": xpu_url}148+ register_data["medium_endpoints"] = {"npu": npu_url}
145 if TENANT_ID != "default":149 if TENANT_ID != "default":
146 register_data["tenant_id"] = TENANT_ID150 register_data["tenant_id"] = TENANT_ID
147 if replay_url:151 if replay_url:
@@ -151,7 +155,7 @@ class ConductorApiClient:
151 try:155 try:
152 with SafeHTTPSClient(timeout=15, **client_args) as client:156 with SafeHTTPSClient(timeout=15, **client_args) as client:
153 client.post("/register", register_data)157 client.post("/register", register_data)
154- mode = "ZMQ+HTTP" if xpu_url else "HTTP-only"158+ mode = "ZMQ+HTTP" if npu_url else "HTTP-only"
155 logger.info(159 logger.info(
156 "HBM DP registered (%s): instance=%s dp=%d replay=%s",160 "HBM DP registered (%s): instance=%s dp=%d replay=%s",
157 mode,161 mode,
@@ -218,12 +222,12 @@ class ConductorApiClient:
218 @classmethod222 @classmethod
219 def _build_medium_endpoints(cls, config, ip: str, dp_rank: int) -> dict[str, str]:223 def _build_medium_endpoints(cls, config, ip: str, dp_rank: int) -> dict[str, str]:
220 """Build the medium_endpoints map from per-medium endpoint patterns."""224 """Build the medium_endpoints map from per-medium endpoint patterns."""
221- xpu_url = cls._resolve_endpoint_url(config.xpu_endpoint, ip, dp_rank)225+ npu_url = cls._resolve_endpoint_url(config.npu_endpoint or config.xpu_endpoint, ip, dp_rank)
222 cpu_url = cls._resolve_endpoint_url(config.cpu_endpoint, ip, dp_rank)226 cpu_url = cls._resolve_endpoint_url(config.cpu_endpoint, ip, dp_rank)
223 disk_url = cls._resolve_endpoint_url(config.disk_endpoint, ip, dp_rank)227 disk_url = cls._resolve_endpoint_url(config.disk_endpoint, ip, dp_rank)
224 fallback = cls._resolve_endpoint_url(config.endpoint, ip, dp_rank)228 fallback = cls._resolve_endpoint_url(config.endpoint, ip, dp_rank)
225 return {229 return {
226- "xpu": xpu_url or fallback or "",230+ "npu": npu_url or fallback or "",
227 "cpu": cpu_url or fallback or "",231 "cpu": cpu_url or fallback or "",
228 "disk": disk_url or fallback or "",232 "disk": disk_url or fallback or "",
229 }233 }
@@ -309,7 +313,7 @@ class ConductorApiClient:
309 313 
310 @classmethod314 @classmethod
311 def query_conductor(cls, instances: list[Instance], encoded_ids: list[int]) -> dict[str, Any]:315 def query_conductor(cls, instances: list[Instance], encoded_ids: list[int]) -> dict[str, Any]:
312- """Query KV conductor for prefix cache overlap scores.316+ """Query KV conductor for prefix cache matched blocks.
313 317 
314 Circuit breaker: after ``_QUERY_CB_THRESHOLD`` consecutive failures,318 Circuit breaker: after ``_QUERY_CB_THRESHOLD`` consecutive failures,
315 skip queries for ``_QUERY_CB_COOLDOWN`` seconds.319 skip queries for ``_QUERY_CB_COOLDOWN`` seconds.
@@ -345,7 +349,7 @@ class ConductorApiClient:
345 try:349 try:
346 with SafeHTTPSClient(timeout=3, **client_args) as client:350 with SafeHTTPSClient(timeout=3, **client_args) as client:
347 response = client.post("/query", query_data)351 response = client.post("/query", query_data)
348- cls._log_hit_summary(response, reg.block_size)352+ logger.info("conductor query response: %s", response)
349 cls._query_failures = 0 # reset on success353 cls._query_failures = 0 # reset on success
350 return response354 return response
351 except Exception as e:355 except Exception as e:
@@ -366,69 +370,6 @@ class ConductorApiClient:
366 )370 )
367 return {}371 return {}
368 372 
369- @classmethod
370- def _log_hit_summary(cls, response: dict[str, Any], block_size: int = 128) -> None:
371- """Log a concise per-instance hit summary from the query response."""
372- if not isinstance(response, dict):
373- return
374- for tenant_id, instances in response.items():
375- if not isinstance(instances, dict):
376- continue
377- for inst_id, imd in instances.items():
378- if not isinstance(imd, dict):
379- continue
380- longest = imd.get("longest_matched", 0) # tokens (blocks × block_size)
381- dp = imd.get("DP", {})
382- total_score = imd.get("total_score", 0)
383- 
384- # Aggregate per-DP hit info for the log line.
385- any_hit = False
386- dp_parts = []
387- media_parts = []
388- if isinstance(dp, dict):
389- for rank, v in sorted(dp.items(), key=lambda x: int(x[0])):
390- if isinstance(v, dict):
391- mt = v.get("matched_tokens", 0)
392- s = v.get("total", 0)
393- xpu_blk = v.get("XPU_blk", 0)
394- cpu_blk = v.get("CPU_blk", 0)
395- disk_blk = v.get("DISK_blk", 0)
396- dp_parts.append(f"{rank}:{mt}t/{s}pts")
397- if xpu_blk or cpu_blk or disk_blk:
398- any_hit = True
399- media_fmt = cls._fmt_medium(xpu_blk, cpu_blk, disk_blk, block_size)
400- if media_fmt:
401- media_parts.append(f" conductor media: {tenant_id}/{inst_id} dp={rank} {media_fmt}")
402- else:
403- dp_parts.append(f"{rank}:{v}t")
404- 
405- parts = [
406- f"matched={longest}t",
407- f"score={total_score}",
408- f"DP={{{','.join(dp_parts)}}}",
409- ]
410- logger.info(
411- "conductor hit: %s/%s %s %s",
412- tenant_id,
413- inst_id,
414- "HIT" if any_hit else "MISS",
415- " ".join(parts),
416- )
417- for mp in media_parts:
418- logger.info(mp)
419- 
420- @staticmethod
421- def _fmt_medium(xpu_blk: int, cpu_blk: int, disk_blk: int, block_size: int) -> str:
422- """Format per-medium hit as e.g. ``XPU=768t(6blk) CPU=0t DISK=0t``."""
423- parts = []
424- for label, blk in [("XPU", xpu_blk), ("CPU", cpu_blk), ("DISK", disk_blk)]:
425- if blk:
426- tok = blk * block_size
427- parts.append(f"{label}={tok}t({blk}blk)")
428- else:
429- parts.append(f"{label}=0t")
430- return " ".join(parts)
431- 
432 @classmethod373 @classmethod
433 def _build_register_payload(cls, instance: Instance, endpoint: Endpoint) -> dict[str, Any]:374 def _build_register_payload(cls, instance: Instance, endpoint: Endpoint) -> dict[str, Any]:
434 """Build registration payload using the unified kv_conductor_config config.375 """Build registration payload using the unified kv_conductor_config config.
Mmotor/coordinator/scheduler/policy/kv_cache_affinity.py+5-6
@@ -259,11 +259,10 @@ class KvCacheAffinityPolicy(WorkloadLedgerMixin, BaseSchedulingPolicy):
259 # already excludes headless endpoints / respects enable_multi_endpoints.259 # already excludes headless endpoints / respects enable_multi_endpoints.
260 for ep in instance.get_all_endpoints():260 for ep in instance.get_all_endpoints():
261 matched_raw = dp_map.get(f"{ep.id}", 0)261 matched_raw = dp_map.get(f"{ep.id}", 0)
262- # Conductor reports per-DP match data. Since the multi-medium scoring262+ # Conductor reports per-DP match data. Since the multi-medium blocks
263- # revision (DpScoring struct), the value is a dict with a "matched_tokens"263+ # revision (DpBlocks struct), the value is a dict with a "matched_tokens"
264- # key (plus "XPU"/"CPU"/"DISK" per-medium breakdown for future264+ # key (plus "npu_blocks"/"cpu_blocks"/"disk_blocks" per-medium block counts);
265- # medium-aware scheduling); older conductors returned a plain int.265+ # older conductors returned a plain int. Handle both.
266- # Handle both.
267 if isinstance(matched_raw, dict):266 if isinstance(matched_raw, dict):
268 matched = matched_raw.get("matched_tokens", 0)267 matched = matched_raw.get("matched_tokens", 0)
269 else:268 else:
@@ -467,7 +466,7 @@ class TokenizerManager(ThreadSafeSingleton):
467 is_vllm_engine = engine_type == "vllm"466 is_vllm_engine = engine_type == "vllm"
468 467 
469 if is_vllm_engine and self._is_deepseek_v4_model(model_path):468 if is_vllm_engine and self._is_deepseek_v4_model(model_path):
470- from vllm.tokenizers.deepseek_v4 import DeepseekV4Tokenizer469+ from vllm.tokenizers.deepseek_v4 import DeepseekV4Tokenizer # pylint: disable=no-name-in-module
471 470 
472 self.tokenizer = DeepseekV4Tokenizer.from_pretrained(model_path, trust_remote_code=True)471 self.tokenizer = DeepseekV4Tokenizer.from_pretrained(model_path, trust_remote_code=True)
473 self._is_dsv4 = True472 self._is_dsv4 = True
Mmotor/kv_conductor/Cargo.toml+3-1
@@ -14,7 +14,9 @@ description = "KV Conductor - Radix tree based KV cache indexer for MindIE-PyMot
14version = "0.1.0"14version = "0.1.0"
15edition = "2021"15edition = "2021"
16authors = ["MindIE Team"]16authors = ["MindIE Team"]
17-license = "Mulan PSL v2"17+# Package is Mulan PSL v2; several sources are Apache-2.0 Derivative Works of
18+# NVIDIA Dynamo kv-router — see THIRD_PARTY_NOTICES.md and licenses/Apache-2.0.txt.
19+license = "MulanPSL-2.0"
18 20 
19[dependencies]21[dependencies]
20axum = "0.7"22axum = "0.7"
Mmotor/kv_conductor/__init__.py+0-15
@@ -74,9 +74,6 @@ def is_available() -> bool:
74def start(74def start(
75 host: str = "::",75 host: str = "::",
76 port: int = 13333,76 port: int = 13333,
77- hbm_weight: int = 3,
78- cpu_weight: int = 2,
79- disk_weight: int = 1,
80 extra_args: list[str] | None = None,77 extra_args: list[str] | None = None,
81) -> subprocess.Popen:78) -> subprocess.Popen:
82 """Launch kv-conductor as a subprocess.79 """Launch kv-conductor as a subprocess.
@@ -87,12 +84,6 @@ def start(
87 Bind address (default ``::`` — dual-stack IPv4/IPv6).84 Bind address (default ``::`` — dual-stack IPv4/IPv6).
88 port:85 port:
89 HTTP listen port.86 HTTP listen port.
90- hbm_weight:
91- Score per matched HBM/XPU block (scoring config).
92- cpu_weight:
93- Score per matched CPU block.
94- disk_weight:
95- Score per matched disk block.
96 extra_args:87 extra_args:
97 Additional CLI arguments forwarded to the binary.88 Additional CLI arguments forwarded to the binary.
98 89 
@@ -118,12 +109,6 @@ def start(
118 host,109 host,
119 "--port",110 "--port",
120 str(port),111 str(port),
121- "--hbm-weight",
122- str(hbm_weight),
123- "--cpu-weight",
124- str(cpu_weight),
125- "--disk-weight",
126- str(disk_weight),
127 ]112 ]
128 if extra_args:113 if extra_args:
129 cmd.extend(extra_args)114 cmd.extend(extra_args)
Mmotor/kv_conductor/src/backend.rs+3-3
@@ -94,7 +94,7 @@ pub enum MatchMode {
94 /// No auto-attach — the subscriber is tied to a fixed dp_rank (YuanRong).94 /// No auto-attach — the subscriber is tied to a fixed dp_rank (YuanRong).
95 None,95 None,
96 /// Match by IP only. An event with `backend_id=<ip>` is applied to96 /// Match by IP only. An event with `backend_id=<ip>` is applied to
97- /// **every** HBM-registered DP whose XPU endpoint IP matches.97+ /// **every** HBM-registered DP whose NPU endpoint IP matches.
98 /// (Mooncake: one master per cluster, backend_id=node IP).98 /// (Mooncake: one master per cluster, backend_id=node IP).
99 IpOnly,99 IpOnly,
100 /// Match by IP **and** dp_rank. An event with `backend_id=<ip>` and100 /// Match by IP **and** dp_rank. An event with `backend_id=<ip>` and
@@ -284,7 +284,7 @@ mod tests {
284 #[test]284 #[test]
285 fn test_resolve_workers_empty_ip_index() {285 fn test_resolve_workers_empty_ip_index() {
286 let index = make_ip_index(vec![]);286 let index = make_ip_index(vec![]);
287- let media = &[StorageMedium::Xpu];287+ let media = &[StorageMedium::Npu];
288 let workers = MatchMode::IpOnly.resolve_workers(Some(&index), "10.0.0.1", 0, media);288 let workers = MatchMode::IpOnly.resolve_workers(Some(&index), "10.0.0.1", 0, media);
289 assert!(workers.is_empty());289 assert!(workers.is_empty());
290 }290 }
@@ -292,7 +292,7 @@ mod tests {
292 #[test]292 #[test]
293 fn test_resolve_workers_multiple_media() {293 fn test_resolve_workers_multiple_media() {
294 let index = make_ip_index(vec![("10.0.0.1", vec![("prefill-0", 0)])]);294 let index = make_ip_index(vec![("10.0.0.1", vec![("prefill-0", 0)])]);
295- let media = &[StorageMedium::Xpu, StorageMedium::Cpu, StorageMedium::Disk];295+ let media = &[StorageMedium::Npu, StorageMedium::Cpu, StorageMedium::Disk];
296 296 
297 let workers = MatchMode::IpOnly.resolve_workers(Some(&index), "10.0.0.1", 0, media);297 let workers = MatchMode::IpOnly.resolve_workers(Some(&index), "10.0.0.1", 0, media);
298 assert_eq!(workers.len(), 3); // 1 DP × 3 media298 assert_eq!(workers.len(), 3); // 1 DP × 3 media
Mmotor/kv_conductor/src/concurrent_tree.rs+99-42
@@ -1,15 +1,30 @@
1-// Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.1+// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2-// MindIE is licensed under Mulan PSL v2.2+// SPDX-FileCopyrightText: Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.
3-// You can use this software according to the terms and conditions of the Mulan PSL v2.3+// SPDX-License-Identifier: Apache-2.0
4-// You may obtain a copy of Mulan PSL v2 at:4+//
5-// http://license.coscl.org.cn/MulanPSL25+// This file is a Derivative Work of NVIDIA Dynamo kv-router
6-// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,6+// (https://github.com/ai-dynamo/dynamo), originally licensed under the
7-// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,7+// Apache License, Version 2.0. Upstream source path:
8-// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.8+// lib/kv-router/src/indexer/concurrent_radix_tree.rs
9-// See the Mulan PSL v2 for more details.9+//
10+// You may obtain a copy of the Apache License at:
11+// http://www.apache.org/licenses/LICENSE-2.0
12+// Local copy: licenses/Apache-2.0.txt
13+// Attribution: THIRD_PARTY_NOTICES.md
14+//
15+// Modified by Huawei Technologies Co., Ltd. for MindIE-PyMotor KV Conductor
16+// (WorkerKey + Arc<FxHashSet> COW workers; lookup owned by Indexer;
17+// find_matches_detailed / PrefixMatch; sweep_stale_nodes; no CleanupState /
18+// metrics / early_exit path). Huawei modifications are also available under
19+// Mulan PSL v2 (http://license.coscl.org.cn/MulanPSL2). Redistribution of
20+// this file must still comply with Apache License 2.0.
10 21 
11//! Thread-safe Concurrent Radix Tree for KV cache block indexing.22//! Thread-safe Concurrent Radix Tree for KV cache block indexing.
12//!23//!
24+//! Derived from NVIDIA Dynamo kv-router `ConcurrentRadixTree`
25+//! (`lib/kv-router/src/indexer/concurrent_radix_tree.rs`, Apache-2.0). See
26+//! `THIRD_PARTY_NOTICES.md`.
27+//!
13//! Uses `Arc<parking_lot::RwLock<Block>>` per node, enabling:28//! Uses `Arc<parking_lot::RwLock<Block>>` per node, enabling:
14//! - Multiple concurrent `find_matches` (read locks only)29//! - Multiple concurrent `find_matches` (read locks only)
15//! - Exclusive `apply_event` (write locks with hand-over-hand ordering)30//! - Exclusive `apply_event` (write locks with hand-over-hand ordering)
@@ -93,6 +108,16 @@ impl Default for ConcurrentRadixTree {
93 }108 }
94}109}
95 110 
111+/// Per-worker HBM prefix match: depth + last engine sequence hash.
112+///
113+/// ``last_seq_hash`` is the breakpoint used to continue into the CPU
114+/// continuation-edge index.
115+#[derive(Debug, Clone, Copy)]
116+pub struct PrefixMatch {
117+ pub depth: u32,
118+ pub last_seq_hash: Option<SequenceBlockHash>,
119+}
120+ 
96impl ConcurrentRadixTree {121impl ConcurrentRadixTree {
97 /// Create a new empty concurrent radix tree.122 /// Create a new empty concurrent radix tree.
98 pub fn new() -> Self {123 pub fn new() -> Self {
@@ -107,14 +132,28 @@ impl ConcurrentRadixTree {
107 132 
108 /// Find matches for a sequence of `LocalBlockHash` values.133 /// Find matches for a sequence of `LocalBlockHash` values.
109 ///134 ///
110- /// Returns per-worker overlap scores indicating the depth of the longest135+ /// Returns per-worker matched block counts indicating the depth of the
111- /// matching prefix for each worker.136+ /// longest matching prefix for each worker.
112- pub fn find_matches(&self, sequence: &[LocalBlockHash]) -> OverlapScores {137+ pub fn find_matches(&self, sequence: &[LocalBlockHash]) -> OverlapBlocks {
138+ let detailed = self.find_matches_detailed(sequence);
139+ let mut overlap = OverlapBlocks::default();
140+ for (worker, m) in detailed {
141+ overlap.update_blocks(worker, m.depth);
142+ }
143+ overlap
144+ }
145+ 
146+ /// Like [`find_matches`], but also returns the last matched sequence hash
147+ /// so callers can continue into lower-tier indexes.
148+ pub fn find_matches_detailed(
149+ &self,
150+ sequence: &[LocalBlockHash],
151+ ) -> FxHashMap<WorkerKey, PrefixMatch> {
113 let t0 = std::time::Instant::now();152 let t0 = std::time::Instant::now();
114- let mut scores = OverlapScores::default();153+ let mut results: FxHashMap<WorkerKey, PrefixMatch> = FxHashMap::default();
115 154 
116 if sequence.is_empty() {155 if sequence.is_empty() {
117- return scores;156+ return results;
118 }157 }
119 158 
120 // Get first child from root (read lock)159 // Get first child from root (read lock)
@@ -125,18 +164,17 @@ impl ConcurrentRadixTree {
125 164 
126 let Some(first_child) = first_child else {165 let Some(first_child) = first_child else {
127 tracing::trace!(seq_len = sequence.len(), "tree miss at root");166 tracing::trace!(seq_len = sequence.len(), "tree miss at root");
128- return scores;167+ return results;
129 };168 };
130 169 
131 // Initialize active workers from first child.170 // Initialize active workers from first child.
132- // Arc clone is O(1) refcount bump instead of O(n) set clone.
133 let mut active: FxHashSet<WorkerKey> = {171 let mut active: FxHashSet<WorkerKey> = {
134 let guard = first_child.read();172 let guard = first_child.read();
135 (*guard.workers).clone()173 (*guard.workers).clone()
136 };174 };
137 175 
138 if active.is_empty() {176 if active.is_empty() {
139- return scores;177+ return results;
140 }178 }
141 179 
142 let mut current = first_child;180 let mut current = first_child;
@@ -153,6 +191,8 @@ impl ConcurrentRadixTree {
153 break;191 break;
154 };192 };
155 193 
194+ let last_seq = { current.read().block_hash };
195+ 
156 // Short-circuit: if only 1 worker remains, just check whether196 // Short-circuit: if only 1 worker remains, just check whether
157 // this single worker is in the child's set.197 // this single worker is in the child's set.
158 if active.len() == 1 {198 if active.len() == 1 {
jason lyu
jason lyujason lyu19 天前

单 worker 短路分支里 last_seq 取当前 block_hash，但 matched_depth 未包含该块，续接时可能用错 breakpoint。

likedislike
@@ -166,20 +206,31 @@ impl ConcurrentRadixTree {
166 matched_depth += 1;206 matched_depth += 1;
167 continue;207 continue;
168 } else {208 } else {
169- scores.update_score(w, matched_depth);209+ results.insert(
210+ w,
211+ PrefixMatch {
212+ depth: matched_depth,
213+ last_seq_hash: last_seq,
214+ },
215+ );
170 active.clear();216 active.clear();
171 break;217 break;
172 }218 }
173 }219 }
174 220 
175 // Reconcile: remove workers that don't have this child block.221 // Reconcile: remove workers that don't have this child block.
176- // Use retain to avoid a second Vec allocation.
177 let guard = block.read();222 let guard = block.read();
178 active.retain(|w| {223 active.retain(|w| {
179 if guard.workers.contains(w) {224 if guard.workers.contains(w) {
180 true225 true
181 } else {226 } else {
182- scores.update_score(w.clone(), matched_depth);227+ results.insert(
228+ w.clone(),
229+ PrefixMatch {
230+ depth: matched_depth,
231+ last_seq_hash: last_seq,
232+ },
233+ );
183 false234 false
184 }235 }
185 });236 });
@@ -193,19 +244,25 @@ impl ConcurrentRadixTree {
193 matched_depth += 1;244 matched_depth += 1;
194 }245 }
195 246 
196- // Drain surviving workers into scores (avoid clone)247+ let last_seq = { current.read().block_hash };
197 for worker in active.drain() {248 for worker in active.drain() {
198- scores.update_score(worker, matched_depth);249+ results.insert(
250+ worker,
251+ PrefixMatch {
252+ depth: matched_depth,
253+ last_seq_hash: last_seq,
254+ },
255+ );
199 }256 }
200 257 
201 tracing::debug!(258 tracing::debug!(
202 seq_len = sequence.len(),259 seq_len = sequence.len(),
203 depth = matched_depth,260 depth = matched_depth,
204- active_workers = scores.scores.len(),261+ active_workers = results.len(),
205 elapsed_us = t0.elapsed().as_micros(),262 elapsed_us = t0.elapsed().as_micros(),
206 "find_matches"263 "find_matches"
207 );264 );
208- scores265+ results
209 }266 }
210 267 
211 // -----------------------------------------------------------------------268 // -----------------------------------------------------------------------
@@ -355,7 +412,7 @@ impl ConcurrentRadixTree {
355 412 
356 /// Remove a worker entirely, cleaning up all tree references.413 /// Remove a worker entirely, cleaning up all tree references.
357 pub fn remove_worker(&self, worker: &WorkerKey, lookup: &mut WorkerLookup) {414 pub fn remove_worker(&self, worker: &WorkerKey, lookup: &mut WorkerLookup) {
358- for (_, block) in lookup.iter() {415+ for block in lookup.values() {
359 block.write().drop_worker(worker);416 block.write().drop_worker(worker);
360 }417 }
361 lookup.clear();418 lookup.clear();
@@ -411,7 +468,7 @@ mod tests {
411 instance_id: instance_id.to_string(),468 instance_id: instance_id.to_string(),
412 backend_id: instance_id.to_string(),469 backend_id: instance_id.to_string(),
413 dp_rank,470 dp_rank,
414- medium: StorageMedium::Xpu,471+ medium: StorageMedium::Npu,
415 }472 }
416 }473 }
417 474 
@@ -432,8 +489,8 @@ mod tests {
432 #[test]489 #[test]
433 fn test_concurrent_find_matches_empty() {490 fn test_concurrent_find_matches_empty() {
434 let tree = ConcurrentRadixTree::new();491 let tree = ConcurrentRadixTree::new();
435- let scores = tree.find_matches(&[LocalBlockHash(1)]);492+ let overlap = tree.find_matches(&[LocalBlockHash(1)]);
436- assert!(scores.is_empty());493+ assert!(overlap.is_empty());
437 }494 }
438 495 
439 #[test]496 #[test]
@@ -445,8 +502,8 @@ mod tests {
445 tree.apply_store(&w1, &mut lookup, &make_store(None, vec![(100, 1)]))502 tree.apply_store(&w1, &mut lookup, &make_store(None, vec![(100, 1)]))
446 .unwrap();503 .unwrap();
447 504 
448- let scores = tree.find_matches(&[LocalBlockHash(1)]);505+ let overlap = tree.find_matches(&[LocalBlockHash(1)]);
449- assert_eq!(scores.scores.get(&w1).copied(), Some(1));506+ assert_eq!(overlap.blocks.get(&w1).copied(), Some(1));
450 }507 }
451 508 
452 #[test]509 #[test]
@@ -462,8 +519,8 @@ mod tests {
462 tree.apply_store(&w1, &mut lookup, &make_store(Some(200), vec![(300, 3)]))519 tree.apply_store(&w1, &mut lookup, &make_store(Some(200), vec![(300, 3)]))
463 .unwrap();520 .unwrap();
464 521 
465- let scores = tree.find_matches(&[LocalBlockHash(1), LocalBlockHash(2), LocalBlockHash(3)]);522+ let overlap = tree.find_matches(&[LocalBlockHash(1), LocalBlockHash(2), LocalBlockHash(3)]);
466- assert_eq!(scores.scores.get(&w1).copied(), Some(3));523+ assert_eq!(overlap.blocks.get(&w1).copied(), Some(3));
467 }524 }
468 525 
469 #[test]526 #[test]
@@ -485,9 +542,9 @@ mod tests {
485 542 
486 assert_eq!(lookup.len(), 1);543 assert_eq!(lookup.len(), 1);
487 544 
488- let scores = tree.find_matches(&[LocalBlockHash(1), LocalBlockHash(2)]);545+ let overlap = tree.find_matches(&[LocalBlockHash(1), LocalBlockHash(2)]);
489 // Only first block matches since second was removed546 // Only first block matches since second was removed
490- assert_eq!(scores.scores.get(&w1).copied(), Some(1));547+ assert_eq!(overlap.blocks.get(&w1).copied(), Some(1));
491 }548 }
492 549 
493 #[test]550 #[test]
@@ -506,8 +563,8 @@ mod tests {
506 tree.remove_worker(&w1, &mut lookup);563 tree.remove_worker(&w1, &mut lookup);
507 564 
508 assert!(lookup.is_empty());565 assert!(lookup.is_empty());
509- let scores = tree.find_matches(&[LocalBlockHash(1)]);566+ let overlap = tree.find_matches(&[LocalBlockHash(1)]);
510- assert!(scores.is_empty());567+ assert!(overlap.is_empty());
511 }568 }
512 569 
513 #[test]570 #[test]
@@ -531,13 +588,13 @@ mod tests {
531 .unwrap();588 .unwrap();
532 589 
533 // Both match first block590 // Both match first block
534- let scores = tree.find_matches(&[LocalBlockHash(1)]);591+ let overlap = tree.find_matches(&[LocalBlockHash(1)]);
535- assert_eq!(scores.scores.get(&w1).copied(), Some(1));592+ assert_eq!(overlap.blocks.get(&w1).copied(), Some(1));
536- assert_eq!(scores.scores.get(&w2).copied(), Some(1));593+ assert_eq!(overlap.blocks.get(&w2).copied(), Some(1));
537 594 
538 // Path (1, 2) matches W1 deeper595 // Path (1, 2) matches W1 deeper
539- let scores = tree.find_matches(&[LocalBlockHash(1), LocalBlockHash(2)]);596+ let overlap = tree.find_matches(&[LocalBlockHash(1), LocalBlockHash(2)]);
540- assert_eq!(scores.scores.get(&w1).copied(), Some(2));597+ assert_eq!(overlap.blocks.get(&w1).copied(), Some(2));
541- assert_eq!(scores.scores.get(&w2).copied(), Some(1));598+ assert_eq!(overlap.blocks.get(&w2).copied(), Some(1));
542 }599 }
543}600}
Mmotor/kv_conductor/src/events/flex_hash.rs+1-1
@@ -11,7 +11,7 @@
11//! `FlexHash` — polymorphic u64 deserialization for msgpack values.11//! `FlexHash` — polymorphic u64 deserialization for msgpack values.
12//!12//!
13//! Supports multiple representations:13//! Supports multiple representations:
14-//! - integer (u64, i64, u32, …)14+//! - non-negative integer (u64, i64, u32, …; negative i64 is rejected)
15//! - decimal string "12345678901234567890"15//! - decimal string "12345678901234567890"
16//! - hex string "0xABCD1234…" or "ABCD1234…"16//! - hex string "0xABCD1234…" or "ABCD1234…"
17//! - binary bytes up to 8 bytes, big-endian (vLLM BlockHash compat)17//! - binary bytes up to 8 bytes, big-endian (vLLM BlockHash compat)
Mmotor/kv_conductor/src/events/mod.rs+18-11
@@ -14,21 +14,28 @@
14//!14//!
15//! 1. **Pool backend** (Mooncake / Memcache / YuanRong): `PoolEvent` with15//! 1. **Pool backend** (Mooncake / Memcache / YuanRong): `PoolEvent` with
16//! `seq_hashes`/`block_hashes`, processed by `apply_pool_event`. No16//! `seq_hashes`/`block_hashes`, processed by `apply_pool_event`. No
17-//! `token_ids` — `tokens_hash` is set equal to `block_hash`.17+//! `token_ids` on the wire — `tokens_hash` comes from the engine offload
18+//! cache, retained content (TTL-bounded),
19+//! or a prior lower-tier insert via the two-phase match (CPU→Disk
20+//! promotion).
18//!21//!
19//! 2. **vLLM engine** (native): `VllmEventMap` with `token_ids` + `block_size`,22//! 2. **vLLM engine** (native): `VllmEventMap` with `token_ids` + `block_size`,
20//! processed by `apply_vllm_event`. `tokens_hash` is re-computed from23//! processed by `apply_vllm_event`. `tokens_hash` is re-computed from
21-//! `token_ids` via `compute_block_hash_for_seq` (XXH3, seed 1337), while24+//! `token_ids` via `compute_block_hash_for_seq` (XXH3, seed 1337); the
22-//! the engine's chained `block_hashes` are kept as25+//! engine's `block_hashes` are matched against `pending_pool` and cached
23-//! `ExternalSequenceBlockHash` for reverse-lookup on `BlockRemoved`.26+//! as `SequenceBlockHash` until the pool confirms (reverse-lookup on
27+//! `BlockRemoved` goes through `evict_pending_blocks`).
24//!28//!
25//! ## Event filtering29//! ## Event filtering
26//!30//!
27-//! Following Dynamo kv-router's approach, events from non-main attention31+//! Events from non-main attention groups (SWA, Mamba, ChunkedLocal, etc.)
28-//! groups (SWA, Mamba, ChunkedLocal, etc.) are filtered out. Only32+//! are filtered out. Only `FullAttention`, `MlaAttention`, and
29-//! `FullAttention`, `MlaAttention`, and `SinkFullAttention` events are33+//! `SinkFullAttention` events are processed (PascalCase or snake_case wire
30-//! processed. This ensures all ingested events share the same `block_size`,34+//! forms such as `mla_attention`); kinds like `sliding_window_mla` are
31-//! avoiding the multi-group hash granularity mismatch problem.35+//! denied. The allow/deny kind set is derived from NVIDIA Dynamo
36+//! kv-router (Apache-2.0); see `THIRD_PARTY_NOTICES.md`. This ensures all
37+//! ingested events share the same `block_size`, avoiding multi-group hash
38+//! granularity mismatch.
32 39 
33pub(crate) mod flex_hash;40pub(crate) mod flex_hash;
34mod helpers;41mod helpers;
@@ -54,8 +61,8 @@ pub(crate) use crate::error::KvConductorError;
54pub(crate) use crate::hashing::compute_block_hash_for_seq;61pub(crate) use crate::hashing::compute_block_hash_for_seq;
55#[allow(unused_imports)]62#[allow(unused_imports)]
56pub(crate) use crate::protocols::{63pub(crate) use crate::protocols::{
57- HbmIpIndex, KvCacheEventData, KvCacheStoreData, KvCacheStoredBlockData, ScoringConfig,64+ HbmIpIndex, KvCacheEventData, KvCacheStoreData, KvCacheStoredBlockData, SequenceBlockHash,
58- SequenceBlockHash, StorageMedium, WorkerKey,65+ StorageMedium, WorkerKey,
59};66};
60 67 
61#[cfg(test)]68#[cfg(test)]
Mmotor/kv_conductor/src/events/pool.rs+18-9
@@ -12,7 +12,9 @@
12//!12//!
13//! Used by all pool backends (Mooncake, Memcache, YuanRong) — the pool13//! Used by all pool backends (Mooncake, Memcache, YuanRong) — the pool
14//! daemon broadcasts KV cache events that this module deserializes and14//! daemon broadcasts KV cache events that this module deserializes and
15-//! matches against offload cache entries.15+//! resolves against the offload cache, retained content, or an
16+//! already-indexed lower tier (two-phase match, pool-first arrivals wait
17+//! in `pending_pool`).
16 18 
17use serde::Deserialize;19use serde::Deserialize;
18 20 
@@ -104,7 +106,7 @@ pub(crate) fn apply_pool_event(
104 // For MatchMode::None (YuanRong), use the subscriber's backend_id which IS106 // For MatchMode::None (YuanRong), use the subscriber's backend_id which IS
105 // the engine instance_id from registration. The event's backend_id may be107 // the engine instance_id from registration. The event's backend_id may be
106 // the pool daemon's IP:port — using it directly would create a different108 // the pool daemon's IP:port — using it directly would create a different
107- // instance_id than the HBM blocks, breaking cross-media score aggregation.109+ // instance_id than the HBM blocks, breaking cross-media block aggregation.
108 // For pool backends (Mooncake/Memcache), the event's backend_id is the node110 // For pool backends (Mooncake/Memcache), the event's backend_id is the node
109 // IP, needed by resolve_workers for hbm_ip_index lookup.111 // IP, needed by resolve_workers for hbm_ip_index lookup.
110 let event_be_id = pool_event.backend_id.as_deref().unwrap_or(backend_id);112 let event_be_id = pool_event.backend_id.as_deref().unwrap_or(backend_id);
@@ -165,14 +167,21 @@ pub(crate) fn apply_pool_event(
165 medium = %worker.medium.as_str(),167 medium = %worker.medium.as_str(),
166 ?preview,168 ?preview,
167 worker = %worker.instance_id,169 worker = %worker.instance_id,
168- "pool confirm: matched cached offload blocks, inserting into tree"170+ "pool confirm: resolved mapping, inserting into tree"
169 );171 );
170- let store_data = KvCacheStoreData {172+ // Apply one `Stored` event per block, each with its own
jason lyu
jason lyujason lyu19 天前

每个 block 单独 apply_event 会加锁多次，对大批量确认可能性能差，考虑批量带 parent_hash。

likedislike
171- parent_hash: None,173+ // `parent_hash` (resolved from the offload cache, retained
172- start_position: None,174+ // content, or a lower-tier lookup), instead of batching them
173- blocks,175+ // into a single event with `parent_hash: None` — batching
174- };176+ // would silently drop continuation-edge chaining.
175- entry.apply_event(worker, &KvCacheEventData::Stored(store_data))?;177+ for (parent_hash, block) in blocks {
178+ let store_data = KvCacheStoreData {
179+ parent_hash,
180+ start_position: None,
181+ blocks: vec![block],
182+ };
183+ entry.apply_event(worker, &KvCacheEventData::Stored(store_data))?;
184+ }
176 }185 }
177 } else if is_removed {186 } else if is_removed {
178 let block_hashes: Vec<u64> = seq_hashes.iter().map(|h| h.0).collect();187 let block_hashes: Vec<u64> = seq_hashes.iter().map(|h| h.0).collect();
Mmotor/kv_conductor/src/events/tests.rs+63-51
@@ -381,7 +381,7 @@ fn test_apply_vllm_block_stored_computes_tokens_hash() {
381 use crate::hashing::compute_block_hash_for_seq;381 use crate::hashing::compute_block_hash_for_seq;
382 use crate::indexer::Indexer;382 use crate::indexer::Indexer;
383 383 
384- let indexer = Indexer::new(ScoringConfig::default());384+ let indexer = Indexer::new();
385 let token_ids = vec![1i64, 2, 3, 4, 5, 6, 7, 8];385 let token_ids = vec![1i64, 2, 3, 4, 5, 6, 7, 8];
386 let block_size = 4u32;386 let block_size = 4u32;
387 387 
@@ -404,7 +404,7 @@ fn test_apply_vllm_block_stored_computes_tokens_hash() {
404 "test-tenant",404 "test-tenant",
405 "test-backend",405 "test-backend",
406 0,406 0,
407- &[StorageMedium::Xpu],407+ &[StorageMedium::Npu],
408 MatchMode::None,408 MatchMode::None,
409 &None,409 &None,
410 block_size,410 block_size,
@@ -420,7 +420,7 @@ fn test_apply_vllm_block_stored_computes_tokens_hash() {
420 instance_id: "test-backend".into(),420 instance_id: "test-backend".into(),
421 backend_id: "test-backend".into(),421 backend_id: "test-backend".into(),
422 dp_rank: 0,422 dp_rank: 0,
423- medium: StorageMedium::Xpu,423+ medium: StorageMedium::Npu,
424 };424 };
425 let lookup = lookups.get(&wk).expect("worker should exist");425 let lookup = lookups.get(&wk).expect("worker should exist");
426 // 2 SHA256 hashes → 2 lookup entries426 // 2 SHA256 hashes → 2 lookup entries
@@ -434,9 +434,9 @@ fn test_apply_vllm_block_stored_computes_tokens_hash() {
434 }434 }
435 435 
436 // Query via find_matches should match (tokens_hash == query hash)436 // Query via find_matches should match (tokens_hash == query hash)
437- let scores = entry.find_matches(&token_ids, block_size, &ScoringConfig::default());437+ let scores = entry.find_matches(&token_ids, block_size);
438 assert!(438 assert!(
439- !scores.scores.is_empty(),439+ !scores.blocks.is_empty(),
440 "query should match stored blocks"440 "query should match stored blocks"
441 );441 );
442}442}
@@ -449,7 +449,7 @@ fn test_apply_vllm_block_stored_computes_tokens_hash() {
449fn test_non_hbm_event_cached_not_in_tree() {449fn test_non_hbm_event_cached_not_in_tree() {
450 use crate::indexer::Indexer;450 use crate::indexer::Indexer;
451 451 
452- let indexer = Indexer::new(ScoringConfig::default());452+ let indexer = Indexer::new();
453 let token_ids = vec![1i64, 2, 3, 4];453 let token_ids = vec![1i64, 2, 3, 4];
454 let block_size = 4u32;454 let block_size = 4u32;
455 455 
@@ -485,9 +485,9 @@ fn test_non_hbm_event_cached_not_in_tree() {
485 }485 }
486 486 
487 // Tree should NOT have the block (not inserted for non-HBM).487 // Tree should NOT have the block (not inserted for non-HBM).
488- let scores = entry.find_matches(&token_ids, block_size, &ScoringConfig::default());488+ let scores = entry.find_matches(&token_ids, block_size);
489 assert!(489 assert!(
490- scores.scores.is_empty(),490+ scores.blocks.is_empty(),
491 "non-HBM events should not be inserted into tree"491 "non-HBM events should not be inserted into tree"
492 );492 );
493}493}
@@ -496,7 +496,7 @@ fn test_non_hbm_event_cached_not_in_tree() {
496fn test_pool_backend_store_matches_cached_block() {496fn test_pool_backend_store_matches_cached_block() {
497 use crate::indexer::Indexer;497 use crate::indexer::Indexer;
498 498 
499- let indexer = Indexer::new(ScoringConfig::default());499+ let indexer = Indexer::new();
500 let token_ids = vec![1i64, 2, 3, 4];500 let token_ids = vec![1i64, 2, 3, 4];
501 let block_size = 4u32;501 let block_size = 4u32;
502 502 
@@ -555,7 +555,7 @@ fn test_pool_backend_store_matches_cached_block() {
555 555 
556 // Tree should now have the block at the pool backend's worker key.556 // Tree should now have the block at the pool backend's worker key.
557 let entry = indexer.get_or_create("test-model", "test-tenant");557 let entry = indexer.get_or_create("test-model", "test-tenant");
558- let scores = entry.find_matches(&token_ids, block_size, &ScoringConfig::default());558+ let scores = entry.find_matches(&token_ids, block_size);
559 // The pool backend worker ("test-pool") should have a match.559 // The pool backend worker ("test-pool") should have a match.
560 let pool_worker = WorkerKey {560 let pool_worker = WorkerKey {
561 instance_id: "test-pool".into(),561 instance_id: "test-pool".into(),
@@ -564,7 +564,7 @@ fn test_pool_backend_store_matches_cached_block() {
564 medium: StorageMedium::Cpu,564 medium: StorageMedium::Cpu,
565 };565 };
566 assert!(566 assert!(
567- scores.scores.contains_key(&pool_worker),567+ scores.blocks.contains_key(&pool_worker),
568 "pool backend store should insert cached block into tree at pool worker"568 "pool backend store should insert cached block into tree at pool worker"
569 );569 );
570}570}
@@ -573,7 +573,7 @@ fn test_pool_backend_store_matches_cached_block() {
573fn test_pool_backend_store_ignores_unknown_hash() {573fn test_pool_backend_store_ignores_unknown_hash() {
574 use crate::indexer::Indexer;574 use crate::indexer::Indexer;
575 575 
576- let indexer = Indexer::new(ScoringConfig::default());576+ let indexer = Indexer::new();
577 let entry = indexer.get_or_create("test-model", "test-tenant");577 let entry = indexer.get_or_create("test-model", "test-tenant");
578 578 
579 // Pool backend stores a block we never cached — now queued in579 // Pool backend stores a block we never cached — now queued in
@@ -628,7 +628,7 @@ fn test_pool_backend_store_ignores_unknown_hash() {
628fn test_pool_backend_remove_evicts_cache() {628fn test_pool_backend_remove_evicts_cache() {
629 use crate::indexer::Indexer;629 use crate::indexer::Indexer;
630 630 
631- let indexer = Indexer::new(ScoringConfig::default());631+ let indexer = Indexer::new();
632 let token_ids = vec![1i64, 2, 3, 4, 5, 6, 7, 8];632 let token_ids = vec![1i64, 2, 3, 4, 5, 6, 7, 8];
633 let block_size = 4u32;633 let block_size = 4u32;
634 634 
@@ -727,16 +727,25 @@ fn test_pool_backend_remove_evicts_cache() {
727 assert!(state.pending_pool.is_empty());727 assert!(state.pending_pool.is_empty());
728 }728 }
729 729 
730- // Tree should still have the pool worker's block for 0xBBB (only 0xAAA was removed).730+ // After removing 0xAAA, 0xBBB remains in the CPU tier. Contiguous prefix
731- let scores = entry.find_matches(&token_ids, block_size, &ScoringConfig::default());731+ // lookup cannot walk past the hole, so find_matches on the full sequence
732+ // reports no hit — verify tier membership directly instead.
732 let pool_worker = WorkerKey {733 let pool_worker = WorkerKey {
733 instance_id: "test-pool".into(),734 instance_id: "test-pool".into(),
734 backend_id: "test-pool".into(),735 backend_id: "test-pool".into(),
735 dp_rank: 0,736 dp_rank: 0,
736 medium: StorageMedium::Cpu,737 medium: StorageMedium::Cpu,
737 };738 };
738- // After removing 0xAAA, only the 0xBBB block remains → still 1 match.739+ assert!(!entry.cpu_tiers.contains_block(0xAAA));
739- assert!(scores.scores.contains_key(&pool_worker));740+ assert!(
741+ entry.cpu_tiers.contains_block(0xBBB),
742+ "removing 0xAAA must not evict the remaining block 0xBBB"
743+ );
744+ assert_eq!(entry.cpu_tiers.worker_block_count(&pool_worker), 1);
745+ assert!(
746+ entry.find_matches(&token_ids, block_size).blocks.is_empty(),
747+ "broken prefix chain should not produce a contiguous match"
748+ );
740}749}
741 750 
742// -----------------------------------------------------------------------751// -----------------------------------------------------------------------
@@ -750,7 +759,7 @@ fn test_pool_backend_remove_evicts_cache() {
750fn test_pool_arrives_before_offload() {759fn test_pool_arrives_before_offload() {
751 use crate::indexer::Indexer;760 use crate::indexer::Indexer;
752 761 
753- let indexer = Indexer::new(ScoringConfig::default());762+ let indexer = Indexer::new();
754 let token_ids = vec![1i64, 2, 3, 4];763 let token_ids = vec![1i64, 2, 3, 4];
755 let block_size = 4u32;764 let block_size = 4u32;
756 765 
@@ -793,8 +802,8 @@ fn test_pool_arrives_before_offload() {
793 }802 }
794 803 
795 // Tree still empty — no offload has arrived yet.804 // Tree still empty — no offload has arrived yet.
796- let scores = entry.find_matches(&token_ids, block_size, &ScoringConfig::default());805+ let scores = entry.find_matches(&token_ids, block_size);
797- assert!(scores.scores.is_empty());806+ assert!(scores.blocks.is_empty());
798 807 
799 // Phase 2: engine offloads to CPU (arrives SECOND).808 // Phase 2: engine offloads to CPU (arrives SECOND).
800 let engine_event = VllmEvent::BlockStored {809 let engine_event = VllmEvent::BlockStored {
@@ -827,7 +836,7 @@ fn test_pool_arrives_before_offload() {
827 }836 }
828 837 
829 // Tree should now have the block at the pool backend's worker key.838 // Tree should now have the block at the pool backend's worker key.
830- let scores = entry.find_matches(&token_ids, block_size, &ScoringConfig::default());839+ let scores = entry.find_matches(&token_ids, block_size);
831 let pool_worker = WorkerKey {840 let pool_worker = WorkerKey {
832 instance_id: "test-pool".into(),841 instance_id: "test-pool".into(),
833 backend_id: "test-pool".into(),842 backend_id: "test-pool".into(),
@@ -835,7 +844,7 @@ fn test_pool_arrives_before_offload() {
835 medium: StorageMedium::Cpu,844 medium: StorageMedium::Cpu,
836 };845 };
837 assert!(846 assert!(
838- scores.scores.contains_key(&pool_worker),847+ scores.blocks.contains_key(&pool_worker),
839 "pool-first ordering: block should be inserted into tree after offload arrives"848 "pool-first ordering: block should be inserted into tree after offload arrives"
840 );849 );
841}850}
@@ -846,7 +855,7 @@ fn test_pool_arrives_before_offload() {
846fn test_pool_arrives_before_offload_multi_worker() {855fn test_pool_arrives_before_offload_multi_worker() {
847 use crate::indexer::Indexer;856 use crate::indexer::Indexer;
848 857 
849- let indexer = Indexer::new(ScoringConfig::default());858+ let indexer = Indexer::new();
850 let token_ids = vec![10i64, 20, 30, 40];859 let token_ids = vec![10i64, 20, 30, 40];
851 let block_size = 4u32;860 let block_size = 4u32;
852 861 
@@ -919,7 +928,7 @@ fn test_pool_arrives_before_offload_multi_worker() {
919 }928 }
920 929 
921 // Both workers have the block in tree.930 // Both workers have the block in tree.
922- let scores = entry.find_matches(&token_ids, block_size, &ScoringConfig::default());931+ let scores = entry.find_matches(&token_ids, block_size);
923 let w0 = WorkerKey {932 let w0 = WorkerKey {
924 instance_id: "pool".into(),933 instance_id: "pool".into(),
925 backend_id: "pool".into(),934 backend_id: "pool".into(),
@@ -932,8 +941,8 @@ fn test_pool_arrives_before_offload_multi_worker() {
932 dp_rank: 1,941 dp_rank: 1,
933 medium: StorageMedium::Cpu,942 medium: StorageMedium::Cpu,
934 };943 };
935- assert!(scores.scores.contains_key(&w0));944+ assert!(scores.blocks.contains_key(&w0));
936- assert!(scores.scores.contains_key(&w1));945+ assert!(scores.blocks.contains_key(&w1));
937}946}
938 947 
939/// Pool event is queued, then a pool removal arrives — pending entry948/// Pool event is queued, then a pool removal arrives — pending entry
@@ -942,7 +951,7 @@ fn test_pool_arrives_before_offload_multi_worker() {
942fn test_pool_removal_cleans_pending() {951fn test_pool_removal_cleans_pending() {
943 use crate::indexer::Indexer;952 use crate::indexer::Indexer;
944 953 
945- let indexer = Indexer::new(ScoringConfig::default());954+ let indexer = Indexer::new();
946 let entry = indexer.get_or_create("m", "t");955 let entry = indexer.get_or_create("m", "t");
947 956 
948 // Pool stored arrives first → queued.957 // Pool stored arrives first → queued.
@@ -1022,7 +1031,7 @@ fn test_pool_removal_cleans_pending() {
1022fn test_offload_then_vllm_removal() {1031fn test_offload_then_vllm_removal() {
1023 use crate::indexer::Indexer;1032 use crate::indexer::Indexer;
1024 1033 
1025- let indexer = Indexer::new(ScoringConfig::default());1034+ let indexer = Indexer::new();
1026 let entry = indexer.get_or_create("m", "t");1035 let entry = indexer.get_or_create("m", "t");
1027 let block_size = 4u32;1036 let block_size = 4u32;
1028 1037 
@@ -1092,7 +1101,7 @@ fn test_offload_then_vllm_removal() {
1092fn test_removal_after_both_matched() {1101fn test_removal_after_both_matched() {
1093 use crate::indexer::Indexer;1102 use crate::indexer::Indexer;
1094 1103 
1095- let indexer = Indexer::new(ScoringConfig::default());1104+ let indexer = Indexer::new();
1096 let token_ids = vec![1i64, 2, 3, 4];1105 let token_ids = vec![1i64, 2, 3, 4];
1097 let block_size = 4u32;1106 let block_size = 4u32;
1098 1107 
@@ -1163,9 +1172,9 @@ fn test_removal_after_both_matched() {
1163 medium: StorageMedium::Cpu,1172 medium: StorageMedium::Cpu,
1164 };1173 };
1165 {1174 {
1166- let scores = entry.find_matches(&token_ids, block_size, &ScoringConfig::default());1175+ let scores = entry.find_matches(&token_ids, block_size);
1167 assert!(1176 assert!(
1168- scores.scores.contains_key(&pool_worker),1177+ scores.blocks.contains_key(&pool_worker),
1169 "after match: block should be in tree"1178 "after match: block should be in tree"
1170 );1179 );
1171 }1180 }
@@ -1199,9 +1208,9 @@ fn test_removal_after_both_matched() {
1199 1208 
1200 // Tree should be empty now (removal succeeded).1209 // Tree should be empty now (removal succeeded).
1201 // Verify via find_matches — no workers should have the block.1210 // Verify via find_matches — no workers should have the block.
1202- let scores_after = entry.find_matches(&token_ids, block_size, &ScoringConfig::default());1211+ let scores_after = entry.find_matches(&token_ids, block_size);
1203 assert!(1212 assert!(
1204- !scores_after.scores.contains_key(&pool_worker),1213+ !scores_after.blocks.contains_key(&pool_worker),
1205 "removal after match: block should be removed from tree"1214 "removal after match: block should be removed from tree"
1206 );1215 );
1207}1216}
@@ -1212,7 +1221,7 @@ fn test_removal_after_both_matched() {
1212fn test_vllm_removal_after_pool_queued() {1221fn test_vllm_removal_after_pool_queued() {
1213 use crate::indexer::Indexer;1222 use crate::indexer::Indexer;
1214 1223 
1215- let indexer = Indexer::new(ScoringConfig::default());1224+ let indexer = Indexer::new();
1216 let entry = indexer.get_or_create("m", "t");1225 let entry = indexer.get_or_create("m", "t");
1217 let block_size = 4u32;1226 let block_size = 4u32;
1218 1227 
@@ -1284,7 +1293,7 @@ fn test_vllm_removal_after_pool_queued() {
1284fn test_duplicate_pool_stored_idempotent() {1293fn test_duplicate_pool_stored_idempotent() {
1285 use crate::indexer::Indexer;1294 use crate::indexer::Indexer;
1286 1295 
1287- let indexer = Indexer::new(ScoringConfig::default());1296+ let indexer = Indexer::new();
1288 let entry = indexer.get_or_create("m", "t");1297 let entry = indexer.get_or_create("m", "t");
1289 1298 
1290 let zmq_event = PoolEvent {1299 let zmq_event = PoolEvent {
@@ -1340,7 +1349,7 @@ fn test_duplicate_pool_stored_idempotent() {
1340fn test_pending_worker_cleanup() {1349fn test_pending_worker_cleanup() {
1341 use crate::indexer::Indexer;1350 use crate::indexer::Indexer;
1342 1351 
1343- let indexer = Indexer::new(ScoringConfig::default());1352+ let indexer = Indexer::new();
1344 let entry = indexer.get_or_create("m", "t");1353 let entry = indexer.get_or_create("m", "t");
1345 1354 
1346 // Queue pool events for two different workers.1355 // Queue pool events for two different workers.
@@ -1394,7 +1403,7 @@ fn test_pending_worker_cleanup() {
1394fn test_cleared_cleans_pending() {1403fn test_cleared_cleans_pending() {
1395 use crate::indexer::Indexer;1404 use crate::indexer::Indexer;
1396 1405 
1397- let indexer = Indexer::new(ScoringConfig::default());1406+ let indexer = Indexer::new();
1398 let entry = indexer.get_or_create("m", "t");1407 let entry = indexer.get_or_create("m", "t");
1399 1408 
1400 // Queue a pool event.1409 // Queue a pool event.
@@ -1465,7 +1474,7 @@ fn test_cleared_cleans_pending() {
1465fn test_sweep_stale_pending() {1474fn test_sweep_stale_pending() {
1466 use crate::indexer::Indexer;1475 use crate::indexer::Indexer;
1467 1476 
1468- let indexer = Indexer::new(ScoringConfig::default());1477+ let indexer = Indexer::new();
1469 let entry = indexer.get_or_create("m", "t");1478 let entry = indexer.get_or_create("m", "t");
1470 1479 
1471 // Queue a pool event.1480 // Queue a pool event.
@@ -1498,7 +1507,7 @@ fn test_sweep_stale_pending() {
1498 assert_eq!(entry.pending_count(), 1);1507 assert_eq!(entry.pending_count(), 1);
1499 1508 
1500 // Sweep with zero TTL → removes everything.1509 // Sweep with zero TTL → removes everything.
1501- let pruned = entry.sweep_stale_pending(std::time::Duration::ZERO);1510+ let pruned = entry.sweep_stale_pending(std::time::Duration::ZERO, std::time::Duration::ZERO);
1502 assert!(pruned > 0, "zero-TTL sweep should remove pending entries");1511 assert!(pruned > 0, "zero-TTL sweep should remove pending entries");
1503 assert_eq!(entry.pending_count(), 0);1512 assert_eq!(entry.pending_count(), 0);
1504}1513}
@@ -1514,7 +1523,7 @@ fn test_vllm_parent_hash_root_level() {
1514 use crate::hashing::compute_block_hash_for_seq;1523 use crate::hashing::compute_block_hash_for_seq;
1515 use crate::indexer::Indexer;1524 use crate::indexer::Indexer;
1516 1525 
1517- let indexer = Indexer::new(ScoringConfig::default());1526+ let indexer = Indexer::new();
1518 let block_size = 4u32;1527 let block_size = 4u32;
1519 let tokens: Vec<i64> = (0..8).collect();1528 let tokens: Vec<i64> = (0..8).collect();
1520 let hashes = compute_block_hash_for_seq(&tokens, block_size);1529 let hashes = compute_block_hash_for_seq(&tokens, block_size);
@@ -1524,9 +1533,9 @@ fn test_vllm_parent_hash_root_level() {
1524 instance_id: "be".into(),1533 instance_id: "be".into(),
1525 backend_id: "be".into(),1534 backend_id: "be".into(),
1526 dp_rank: 0,1535 dp_rank: 0,
1527- medium: StorageMedium::Xpu,1536+ medium: StorageMedium::Npu,
1528 };1537 };
1529- let media = &[StorageMedium::Xpu];1538+ let media = &[StorageMedium::Npu];
1530 1539 
1531 // 2-block event, no parent — these form a root chain internally.1540 // 2-block event, no parent — these form a root chain internally.
1532 apply_vllm_event(1541 apply_vllm_event(
@@ -1551,9 +1560,9 @@ fn test_vllm_parent_hash_root_level() {
1551 .unwrap();1560 .unwrap();
1552 1561 
1553 let entry = indexer.get_or_create("m", "t");1562 let entry = indexer.get_or_create("m", "t");
1554- let scores = entry.find_matches(&tokens, block_size, &ScoringConfig::default());1563+ let scores = entry.find_matches(&tokens, block_size);
1555 assert!(1564 assert!(
1556- scores.scores.contains_key(&wk),1565+ scores.blocks.contains_key(&wk),
1557 "should match at least 1 block"1566 "should match at least 1 block"
1558 );1567 );
1559 1568 
@@ -1570,7 +1579,7 @@ fn test_vllm_parent_hash_cross_event_chain() {
1570 use crate::hashing::compute_block_hash_for_seq;1579 use crate::hashing::compute_block_hash_for_seq;
1571 use crate::indexer::Indexer;1580 use crate::indexer::Indexer;
1572 1581 
1573- let indexer = Indexer::new(ScoringConfig::default());1582+ let indexer = Indexer::new();
1574 let block_size = 4u32;1583 let block_size = 4u32;
1575 let tokens: Vec<i64> = (0..16).collect();1584 let tokens: Vec<i64> = (0..16).collect();
1576 let hashes = compute_block_hash_for_seq(&tokens, block_size);1585 let hashes = compute_block_hash_for_seq(&tokens, block_size);
@@ -1580,9 +1589,9 @@ fn test_vllm_parent_hash_cross_event_chain() {
1580 instance_id: "be".into(),1589 instance_id: "be".into(),
1581 backend_id: "be".into(),1590 backend_id: "be".into(),
1582 dp_rank: 0,1591 dp_rank: 0,
1583- medium: StorageMedium::Xpu,1592+ medium: StorageMedium::Npu,
1584 };1593 };
1585- let media = &[StorageMedium::Xpu];1594+ let media = &[StorageMedium::Npu];
1586 1595 
1587 // Event 0: blocks 0x100, 0x200, tokens[0..8], no parent.1596 // Event 0: blocks 0x100, 0x200, tokens[0..8], no parent.
1588 apply_vllm_event(1597 apply_vllm_event(
@@ -1629,10 +1638,13 @@ fn test_vllm_parent_hash_cross_event_chain() {
1629 .unwrap();1638 .unwrap();
1630 1639 
1631 let entry = indexer.get_or_create("m", "t");1640 let entry = indexer.get_or_create("m", "t");
1632- let scores = entry.find_matches(&tokens, block_size, &ScoringConfig::default());1641+ let scores = entry.find_matches(&tokens, block_size);
1633- let score = scores.scores.get(&wk).expect("should match HBM chain");1642+ let matched = scores.blocks.get(&wk).expect("should match HBM chain");
1634- // 4-block chain → depth=4 × hbm_weight=3 = 121643+ // Per-medium overlap is matched block count (no weighted scoring).
1635- assert_eq!(*score, 12, "depth=4 × 3 = 12, got {score}");1644+ assert_eq!(
1645+ *matched, 4,
1646+ "4-block HBM chain should match depth=4, got {matched}"
1647+ );
1636 1648 
1637 // Verify parent-not-found error.1649 // Verify parent-not-found error.
1638 let result = apply_vllm_event(1650 let result = apply_vllm_event(
Mmotor/kv_conductor/src/events/vllm.rs+73-27
@@ -7,11 +7,20 @@
7// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,7// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
8// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.8// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
9// See the Mulan PSL v2 for more details.9// See the Mulan PSL v2 for more details.
10+//
11+// Portions: the main-attention allow/deny kind set used by
12+// `is_main_attention_kind` is derived from NVIDIA Dynamo kv-router
13+// `lib/kv-router/src/zmq_wire/filter.rs` (Apache-2.0).
14+// Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
15+// See ../THIRD_PARTY_NOTICES.md and ../licenses/Apache-2.0.txt.
16+// The msgspec/JSON visitor and event-application logic in this file are
17+// Huawei original work under Mulan PSL v2.
10 18 
11//! vLLM-native event types, parsing, and application logic.19//! vLLM-native event types, parsing, and application logic.
12//!20//!
13//! Handles the msgspec ``array_like`` wire format with tag-based dispatch,21//! Handles the msgspec ``array_like`` wire format with tag-based dispatch,
14//! attention-group filtering, and two-phase offload/pool insertion.22//! attention-group filtering, and two-phase offload/pool insertion.
23+//! Attention-kind filtering policy: see `THIRD_PARTY_NOTICES.md`.
15 24 
16use serde::Deserialize;25use serde::Deserialize;
17 26 
@@ -363,18 +372,37 @@ fn parse_block_removed_values(
363// ---------------------------------------------------------------------------372// ---------------------------------------------------------------------------
364 373 
365/// Returns `true` if `kind` is a main attention type whose events should be374/// Returns `true` if `kind` is a main attention type whose events should be
366-/// ingested. Following Dynamo kv-router, only `FullAttention`,375+/// ingested. Allow/deny kind names follow NVIDIA Dynamo kv-router
367-/// `MlaAttention`, and `SinkFullAttention` qualify. Events with no376+/// `zmq_wire/filter.rs` (Apache-2.0; see `THIRD_PARTY_NOTICES.md`): only
368-/// `kv_cache_spec_kind` (older vLLM versions) are kept for backward compat.377+/// `FullAttention`, `MlaAttention`, and `SinkFullAttention` qualify.
378+/// Events with no `kv_cache_spec_kind` (older vLLM versions) are kept for
379+/// backward compat.
380+///
381+/// Matching is case-insensitive and ignores underscores so both PascalCase
382+/// (`MlaAttention`) and vLLM wire snake_case (`mla_attention`) are matched;
383+/// denied kinds (e.g. `sliding_window_mla`) are filtered out.
369pub(crate) fn is_main_attention_kind(kind: Option<&str>) -> bool {384pub(crate) fn is_main_attention_kind(kind: Option<&str>) -> bool {
370- match kind {385+ let Some(kind) = kind else {
371- None => true,386+ return true;
372- Some("FullAttention") | Some("MlaAttention") | Some("SinkFullAttention") => true,387+ };
373- Some("SlidingWindow")388+ // Normalize: lowercase + strip underscores → "MlaAttention"/"mla_attention"
374- | Some("Mamba")389+ // both become "mlaattention".
375- | Some("ChunkedLocalAttention")390+ let normalized: String = kind
376- | Some("EncoderOnlyAttention")391+ .chars()
377- | Some("CrossAttention") => false,392+ .filter(|c| *c != '_')
393+ .flat_map(|c| c.to_lowercase())
394+ .collect();
395+ match normalized.as_str() {
396+ "fullattention" | "mlaattention" | "sinkfullattention" => true,
397+ // Non-main groups (SWA / Mamba / local / encoder / cross).
398+ // `slidingwindowmla` covers wire form `sliding_window_mla`.
399+ "slidingwindow"
400+ | "slidingwindowmla"
401+ | "mamba"
402+ | "chunkedlocalattention"
403+ | "encoderonlyattention"
404+ | "crossattention" => false,
405+ // Unknown future kinds — forward compat (same as before).
378 _ => true,406 _ => true,
379 }407 }
380}408}
@@ -525,7 +553,7 @@ pub(crate) fn parse_vllm_batch(payload: &[u8]) -> Option<(Vec<VllmEvent>, u32)>
525/// us to re-compute `tokens_hash` (XXH3 content hash). The behaviour553/// us to re-compute `tokens_hash` (XXH3 content hash). The behaviour
526/// depends on the storage medium:554/// depends on the storage medium:
527///555///
528-/// - **HBM** (XPU/GPU): insert directly into the radix tree.556+/// - **HBM** (NPU): insert directly into the radix tree.
529/// - **Non-HBM** (CPU/DISK): bidirectional matching — cache the557/// - **Non-HBM** (CPU/DISK): bidirectional matching — cache the
530/// `block_hash → tokens_hash` mapping and check for pending pool events558/// `block_hash → tokens_hash` mapping and check for pending pool events
531/// that arrived earlier. If a match is found the block enters the tree559/// that arrived earlier. If a match is found the block enters the tree
@@ -535,6 +563,10 @@ pub(crate) fn parse_vllm_batch(payload: &[u8]) -> Option<(Vec<VllmEvent>, u32)>
535/// the block on a different node than the engine that offloaded it — the563/// the block on a different node than the engine that offloaded it — the
536/// engine's offloading event tells us *what* was offloaded, and the pool564/// engine's offloading event tells us *what* was offloaded, and the pool
537/// backend's event tells us *where* it was placed.565/// backend's event tells us *where* it was placed.
566+///
567+/// After the first pool confirmation, `(tokens_hash, parent_hash)` is
568+/// retained so a later pool medium (e.g. Disk SSD offload) can reuse
569+/// content without another engine event.
538#[allow(clippy::too_many_arguments)]570#[allow(clippy::too_many_arguments)]
539pub(crate) fn apply_vllm_event(571pub(crate) fn apply_vllm_event(
540 indexer: &Indexer,572 indexer: &Indexer,
@@ -570,9 +602,8 @@ pub(crate) fn apply_vllm_event(
570 return Ok(());602 return Ok(());
571 }603 }
572 604 
573- let event_medium = medium.as_deref().unwrap_or("xpu");605+ let event_medium = medium.as_deref().unwrap_or("npu");
574- let is_non_hbm = !event_medium.eq_ignore_ascii_case("xpu")606+ let is_non_hbm = !StorageMedium::is_hbm_key(event_medium);
575- && !event_medium.eq_ignore_ascii_case("gpu");
576 607 
577 let computed_hashes: Vec<u64> = if token_ids.is_empty() || *block_size == 0 {608 let computed_hashes: Vec<u64> = if token_ids.is_empty() || *block_size == 0 {
578 block_hashes.to_vec()609 block_hashes.to_vec()
@@ -590,20 +621,29 @@ pub(crate) fn apply_vllm_event(
590 let entry = indexer.get_or_create(model_name, tenant_id);621 let entry = indexer.get_or_create(model_name, tenant_id);
591 622 
592 if is_non_hbm {623 if is_non_hbm {
593- let pairs: Vec<(u64, u64)> = (0..num)624+ // Walk the offload chain so each block carries its own
594- .map(|i| (block_hashes[i], computed_hashes[i]))625+ // `parent_hash`: the first block's parent is the event's
595- .collect();626+ // `parent_block_hash`, and each subsequent block's parent is
627+ // the immediately preceding block in this same chain. This
628+ // preserves continuation-edge semantics across the two-phase
629+ // offload/pool confirmation protocol.
630+ let mut parent = *parent_block_hash;
631+ let mut triples: Vec<(u64, u64, Option<u64>)> = Vec::with_capacity(num);
632+ for i in 0..num {
633+ triples.push((block_hashes[i], computed_hashes[i], parent));
634+ parent = Some(block_hashes[i]);
635+ }
596 636 
597- let preview_hashes: Vec<u64> = pairs.iter().take(4).map(|p| p.0).collect();637+ let preview_hashes: Vec<u64> = triples.iter().take(4).map(|p| p.0).collect();
598 tracing::trace!(638 tracing::trace!(
599 model = %model_name, tenant = %tenant_id,639 model = %model_name, tenant = %tenant_id,
600- num = pairs.len(),640+ num = triples.len(),
601 ?preview_hashes,641 ?preview_hashes,
602 medium = %event_medium,642 medium = %event_medium,
603 "vLLM non-HBM: ingesting offload blocks"643 "vLLM non-HBM: ingesting offload blocks"
604 );644 );
605 645 
606- let matched = entry.ingest_offload_blocks(&pairs);646+ let matched = entry.ingest_offload_blocks(&triples);
607 let total_matched: usize = matched.values().map(|v| v.len()).sum();647 let total_matched: usize = matched.values().map(|v| v.len()).sum();
608 648 
609 if !matched.is_empty() {649 if !matched.is_empty() {
@@ -614,13 +654,19 @@ pub(crate) fn apply_vllm_event(
614 medium = %event_medium,654 medium = %event_medium,
615 "vLLM non-HBM: matched pending pool events, applying to tree"655 "vLLM non-HBM: matched pending pool events, applying to tree"
616 );656 );
657+ // Apply one `Stored` event per block, each with its own
658+ // `parent_hash`, instead of batching them under
659+ // `parent_hash: None` — batching would silently drop
660+ // continuation-edge chaining between blocks.
617 for (worker, blocks) in matched {661 for (worker, blocks) in matched {
618- let store_data = KvCacheStoreData {662+ for (parent_hash, block) in blocks {
619- parent_hash: None,663+ let store_data = KvCacheStoreData {
620- start_position: None,664+ parent_hash,
621- blocks,665+ start_position: None,
622- };666+ blocks: vec![block],
623- entry.apply_event(&worker, &KvCacheEventData::Stored(store_data))?;667+ };
668+ entry.apply_event(&worker, &KvCacheEventData::Stored(store_data))?;
669+ }
624 }670 }
625 }671 }
626 672 
Mmotor/kv_conductor/src/hashing.rs+24-10
@@ -1,22 +1,36 @@
1-// Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.1+// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2-// MindIE is licensed under Mulan PSL v2.2+// SPDX-FileCopyrightText: Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.
3-// You can use this software according to the terms and conditions of the Mulan PSL v2.3+// SPDX-License-Identifier: Apache-2.0
4-// You may obtain a copy of Mulan PSL v2 at:4+//
5-// http://license.coscl.org.cn/MulanPSL25+// This file is a Derivative Work of NVIDIA Dynamo kv-router hashing helpers
6-// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,6+// in lib/kv-router/src/protocols.rs (XXH3_SEED, compute_block_hash,
7-// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,7+// little-endian xxh3 block hashing), originally licensed under the Apache
8-// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.8+// License, Version 2.0. Upstream project: https://github.com/ai-dynamo/dynamo
9-// See the Mulan PSL v2 for more details.9+//
10+// You may obtain a copy of the Apache License at:
11+// http://www.apache.org/licenses/LICENSE-2.0
12+// Local copy: licenses/Apache-2.0.txt
13+// Attribution: THIRD_PARTY_NOTICES.md
14+//
15+// Modified by Huawei Technologies Co., Ltd. for MindIE-PyMotor KV Conductor
16+// (i64 token input with per-chunk conversion, rayon parallel batching,
17+// include partial trailing block via div_ceil). Huawei modifications are
18+// also available under Mulan PSL v2 (http://license.coscl.org.cn/MulanPSL2).
19+// Redistribution of this file must still comply with Apache License 2.0.
10 20 
11//! XXH3-based token block hashing.21//! XXH3-based token block hashing.
12//!22//!
23+//! Derived from NVIDIA Dynamo kv-router hashing helpers in
24+//! `lib/kv-router/src/protocols.rs` (Apache-2.0). See `THIRD_PARTY_NOTICES.md`.
25+//!
13//! Computes `LocalBlockHash` values from token sequences using a sliding-window26//! Computes `LocalBlockHash` values from token sequences using a sliding-window
14 27 
15use xxhash_rust::xxh3;28use xxhash_rust::xxh3;
16 29 
17use crate::protocols::LocalBlockHash;30use crate::protocols::LocalBlockHash;
18 31 
19-/// Seed for XXH3 hashing, consistent with Dynamo kv-router.32+/// Seed for XXH3 hashing (same value as NVIDIA Dynamo kv-router for wire
33+/// interoperability).
20pub const XXH3_SEED: u64 = 1337;34pub const XXH3_SEED: u64 = 1337;
21 35 
22/// Compute the hash of arbitrary data.36/// Compute the hash of arbitrary data.
Rmotor/kv_conductor/src/indexer.rs→motor/kv_conductor/src/indexer/mod.rs+541-558
@@ -8,19 +8,23 @@
8// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.8// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
9// See the Mulan PSL v2 for more details.9// See the Mulan PSL v2 for more details.
10 10 
11-//! Per-(model, tenant) radix tree indexer with per-medium scoring.11+//! Per-(model, tenant) radix tree indexer with per-medium block matching.
12//!12//!
13//! Each `IndexerEntry` manages:13//! Each `IndexerEntry` manages:
14//!14//!
15-//! - **HBM tree** (`hbm_tree`) — prefix-chain radix tree for XPU/GPU blocks15+//! - **HBM tree** (`hbm_tree`) — prefix-chain radix tree for NPU blocks.
16-//! (weight ×3). Only HBM blocks need prefix-chain matching.16+//! - **CPU / Disk continuation indexes** (`cpu_tiers` / `disk_tiers`) —
17-//! - **CPU flat map** (`cpu_blocks`) — ``tokens_hash → {workers}`` (weight ×2).17+//! ``(parent_seq_hash, tokens_hash) → child`` edges (see `lower_tier`
18-//! - **Disk flat map** (`disk_blocks`) — ``tokens_hash → {workers}`` (weight ×1).18+//! and `THIRD_PARTY_NOTICES.md`). CPU continues from the HBM
19-//! - **non_hbm_cache** — engine offload ``block_hash → tokens_hash`` for19+//! breakpoint; Disk continues from ``max(HBM, CPU)`` (CPU preferred
20-//! two-phase pool confirmation.20+//! when it extends further). Root chains are walked unconditionally so
21+//! longer lower-tier replicas are never hidden by shorter upstream hits.
22+//! - **offload_pool_state** — bidirectional offload/pool event matching
23+//! (see [`OffloadPoolState`]). The `offload` side now also carries the
24+//! originating `parent_hash` so that lower-tier continuation edges are
25+//! correctly chained once the pool backend confirms placement.
21//!26//!
22-//! Scoring: each matched HBM block = 3 pts, CPU block = 2 pts, disk block = 1 pt.27+//! Query results report matched block counts per medium (no weighted scoring).
23-//! A DP's total score is the sum across all its media.
24 28 
25use std::collections::HashMap;29use std::collections::HashMap;
26use std::sync::atomic::{AtomicU64, Ordering};30use std::sync::atomic::{AtomicU64, Ordering};
@@ -30,33 +34,70 @@ use dashmap::DashMap;
30use parking_lot::RwLock;34use parking_lot::RwLock;
31use rustc_hash::FxHashMap;35use rustc_hash::FxHashMap;
32use rustc_hash::FxHashSet;36use rustc_hash::FxHashSet;
37+use serde::Serialize;
33 38 
34-use crate::concurrent_tree::{ConcurrentRadixTree, WorkerLookup};39+use crate::concurrent_tree::{ConcurrentRadixTree, PrefixMatch, WorkerLookup};
35use crate::error::KvConductorError;40use crate::error::KvConductorError;
36use crate::hashing::compute_block_hash_for_seq;41use crate::hashing::compute_block_hash_for_seq;
42+use crate::lower_tier::{ContiguousHit, LowerTierContinuation, LowerTierIndexer};
37use crate::protocols::*;43use crate::protocols::*;
38 44 
39/// Number of HBM removals after which `sweep_stale_nodes` is triggered45/// Number of HBM removals after which `sweep_stale_nodes` is triggered
40/// to reclaim orphan tree nodes.46/// to reclaim orphan tree nodes.
41const HBM_SWEEP_THRESHOLD: u64 = 1000;47const HBM_SWEEP_THRESHOLD: u64 = 1000;
42 48 
43-/// Minimum number of block hashes to trigger parallel CPU/Disk flat lookup.
44-const FLAT_PAR_THRESHOLD: usize = 4096;
45- 
46/// Number of ingest operations after which `sweep_stale_pending` is triggered49/// Number of ingest operations after which `sweep_stale_pending` is triggered
47-/// to evict expired entries from the pending pool cache.50+/// to evict expired `pending_pool` / `content` entries.
48const PENDING_SWEEP_THRESHOLD: u64 = 100;51const PENDING_SWEEP_THRESHOLD: u64 = 100;
49/// TTL for stale pending pool entries (60 seconds).52/// TTL for stale pending pool entries (60 seconds).
50const PENDING_TTL: std::time::Duration = std::time::Duration::from_secs(60);53const PENDING_TTL: std::time::Duration = std::time::Duration::from_secs(60);
54+/// TTL for retained `content` entries. Content must survive the CPU→Disk
55+/// migration window (CPU tier eviction before the Disk store event arrives),
56+/// so it is deliberately longer than [`PENDING_TTL`]; entries are cleared
57+/// once the window closes.
58+const CONTENT_TTL: std::time::Duration = std::time::Duration::from_secs(300);
51 59 
52-/// Per-worker reverse-lookup for non-HBM flat stores.60+/// Upstream-tier match breakpoint used to continue into the next lower tier.
53-/// Maps ``SequenceBlockHash → LocalBlockHash(XXH3 tokens_hash)``.61+#[derive(Debug, Clone)]
54-type FlatLookup = FxHashMap<u64, u64>;62+struct TierBreakpoint {
63+ instance_id: String,
64+ dp_rank: DpRank,
65+ /// Absolute index in the query hash sequence where the next tier starts.
66+ end_pos: usize,
67+ /// Sequence hash of the last matched upstream block.
68+ last_seq: SequenceBlockHash,
69+}
55 70 
56// ---------------------------------------------------------------------------71// ---------------------------------------------------------------------------
57// Two-phase offload/pool matching protocol72// Two-phase offload/pool matching protocol
58// ---------------------------------------------------------------------------73// ---------------------------------------------------------------------------
59 74 
75+/// Cached Phase-1 offload mapping: `tokens_hash` plus the optional
76+/// `parent_hash` (the engine's immediately preceding block in the offload
77+/// chain, or the chain's original `parent_block_hash` for the first block).
78+///
79+/// Carrying `parent_hash` alongside `tokens_hash` allows Phase-2 confirmation
80+/// to insert this block as a continuation edge from the correct predecessor,
81+/// rather than always chaining from root.
82+#[derive(Debug, Clone, Copy)]
83+pub(crate) struct OffloadCacheEntry {
84+ pub(crate) tokens_hash: u64,
85+ pub(crate) parent_hash: Option<u64>,
86+}
87+ 
88+/// Confirmed content kept for a possible later pool medium (Disk promotion),
89+/// with insert time for TTL eviction ([`CONTENT_TTL`]).
90+///
91+/// Unlike the unconfirmed `offload` entries (which have no TTL — bounded by
92+/// the number of outstanding offloaded blocks), content survives lower-tier
93+/// removal so a Disk store event arriving after CPU eviction can still resolve
94+/// the mapping; the TTL bounds how long it lingers.
95+#[derive(Debug, Clone)]
96+pub(crate) struct ContentEntry {
97+ pub(crate) entry: OffloadCacheEntry,
98+ pub(crate) inserted_at: std::time::Instant,
99+}
100+ 
60/// A pool backend event waiting for its corresponding offload event.101/// A pool backend event waiting for its corresponding offload event.
61///102///
62/// Equality and hashing consider **only** the `worker` field — `inserted_at`103/// Equality and hashing consider **only** the `worker` field — `inserted_at`
@@ -84,16 +125,33 @@ impl std::hash::Hash for PendingPoolEvent {
84 125 
85/// Combined state for the two-phase offload/pool matching protocol.126/// Combined state for the two-phase offload/pool matching protocol.
86///127///
87-/// **Invariant**: a `block_hash` exists in **at most one** of the two maps128+/// **Invariant**: a `block_hash` is never in both `offload` and `content`.
88-/// at any time. Once both sides have arrived, the entry is removed from129+/// After the first pool confirmation, the entry moves from `offload` into
89-/// both maps and inserted into the radix tree (or flat store).130+/// `content` — always retained for a possible later Disk promotion (the
131+/// mapping is a copy of the tier data plus a bounded migration-window
132+/// residue; `content` is TTL-swept so memory stays bounded).
133+///
134+/// Lifecycle:
135+/// - `offload`: unconfirmed engine offloads. **No TTL** — pool confirmation
136+/// may arrive arbitrarily late, so entries are bounded only by the number
137+/// of outstanding offloaded blocks (original two-phase design) and cleared
138+/// on match or explicit removal.
139+/// - `content`: confirmed `(tokens_hash, parent_hash)` kept for a later pool
140+/// medium. **TTL-evicted** ([`CONTENT_TTL`]) — it must survive lower-tier
141+/// removal so a Disk store arriving after CPU eviction still resolves, and
142+/// the TTL bounds how long a block that never reaches Disk lingers.
143+/// - `pending_pool`: pool-first arrivals, TTL [`PENDING_TTL`].
90///144///
91/// Uses a single `RwLock` so that cross-cache operations are atomic without145/// Uses a single `RwLock` so that cross-cache operations are atomic without
92/// lock-ordering deadlock risk.146/// lock-ordering deadlock risk.
93#[derive(Debug, Default)]147#[derive(Debug, Default)]
94pub(crate) struct OffloadPoolState {148pub(crate) struct OffloadPoolState {
95- /// `block_hash → tokens_hash`: offload events waiting for pool confirmation.149+ /// `block_hash → OffloadCacheEntry`: offload events waiting for the
96- pub(crate) offload: FxHashMap<u64, u64>,150+ /// **first** pool confirmation. Not TTL-swept (see struct docs).
151+ pub(crate) offload: FxHashMap<u64, OffloadCacheEntry>,
152+ /// `block_hash → ContentEntry`: retained after the first lower-tier
153+ /// insert (always, see struct docs). Swept with [`CONTENT_TTL`].
154+ pub(crate) content: FxHashMap<u64, ContentEntry>,
97 /// `block_hash → workers`: pool events waiting for offload `tokens_hash`.155 /// `block_hash → workers`: pool events waiting for offload `tokens_hash`.
98 /// Values are `FxHashSet` to deduplicate repeated deliveries.156 /// Values are `FxHashSet` to deduplicate repeated deliveries.
99 pub(crate) pending_pool: FxHashMap<u64, FxHashSet<PendingPoolEvent>>,157 pub(crate) pending_pool: FxHashMap<u64, FxHashSet<PendingPoolEvent>>,
@@ -108,20 +166,15 @@ pub struct IndexerKey {
108 166 
109/// An indexer entry for one (model, tenant) pair.167/// An indexer entry for one (model, tenant) pair.
110pub struct IndexerEntry {168pub struct IndexerEntry {
111- /// HBM prefix-chain radix tree (XPU workers, weight ×3).169+ /// HBM prefix-chain radix tree (NPU workers).
112 pub hbm_tree: Arc<ConcurrentRadixTree>,170 pub hbm_tree: Arc<ConcurrentRadixTree>,
113 /// HBM per-worker reverse lookups: WorkerKey → WorkerLookup.171 /// HBM per-worker reverse lookups: WorkerKey → WorkerLookup.
114 pub lookups: Arc<RwLock<FxHashMap<WorkerKey, WorkerLookup>>>,172 pub lookups: Arc<RwLock<FxHashMap<WorkerKey, WorkerLookup>>>,
115 173 
116- /// CPU flat blocks: tokens_hash → set of workers with this block cached.174+ /// CPU continuation-edge index.
117- pub cpu_blocks: Arc<RwLock<FxHashMap<LocalBlockHash, FxHashSet<WorkerKey>>>>,175+ pub cpu_tiers: Arc<LowerTierIndexer>,
118- /// CPU per-worker reverse lookups: WorkerKey → (seq_hash → tokens_hash).176+ /// Disk continuation-edge index.
119- cpu_lookups: Arc<RwLock<FxHashMap<WorkerKey, FlatLookup>>>,177+ pub disk_tiers: Arc<LowerTierIndexer>,
120- 
121- /// Disk flat blocks: tokens_hash → set of workers with this block cached.
122- pub disk_blocks: Arc<RwLock<FxHashMap<LocalBlockHash, FxHashSet<WorkerKey>>>>,
123- /// Disk per-worker reverse lookups.
124- disk_lookups: Arc<RwLock<FxHashMap<WorkerKey, FlatLookup>>>,
125 178 
126 /// Bidirectional offload/pool event matching state.179 /// Bidirectional offload/pool event matching state.
127 /// See [`OffloadPoolState`] for the invariant.180 /// See [`OffloadPoolState`] for the invariant.
@@ -137,150 +190,233 @@ pub struct IndexerEntry {
137 190 
138impl Default for IndexerEntry {191impl Default for IndexerEntry {
139 fn default() -> Self {192 fn default() -> Self {
140- Self {193+ Self::new()
141- hbm_tree: Arc::new(ConcurrentRadixTree::new()),
142- lookups: Arc::new(RwLock::new(FxHashMap::default())),
143- cpu_blocks: Arc::new(RwLock::new(FxHashMap::default())),
144- cpu_lookups: Arc::new(RwLock::new(FxHashMap::default())),
145- disk_blocks: Arc::new(RwLock::new(FxHashMap::default())),
146- disk_lookups: Arc::new(RwLock::new(FxHashMap::default())),
147- offload_pool_state: Arc::new(RwLock::new(OffloadPoolState::default())),
148- hbm_removal_count: AtomicU64::new(0),
149- pending_ingest_count: AtomicU64::new(0),
150- }
151 }194 }
152}195}
153 196 
154impl IndexerEntry {197impl IndexerEntry {
155 pub fn new() -> Self {198 pub fn new() -> Self {
156- Self::default()199+ Self {
200+ hbm_tree: Arc::new(ConcurrentRadixTree::new()),
201+ lookups: Arc::new(RwLock::new(FxHashMap::default())),
202+ cpu_tiers: Arc::new(LowerTierIndexer::new()),
203+ disk_tiers: Arc::new(LowerTierIndexer::new()),
204+ offload_pool_state: Arc::new(RwLock::new(OffloadPoolState::default())),
205+ hbm_removal_count: AtomicU64::new(0),
206+ pending_ingest_count: AtomicU64::new(0),
207+ }
157 }208 }
158 209 
159 // -----------------------------------------------------------------------210 // -----------------------------------------------------------------------
160 // Query211 // Query
161 // -----------------------------------------------------------------------212 // -----------------------------------------------------------------------
162 213 
163- pub fn find_matches(214+ pub fn find_matches(&self, token_ids: &[i64], block_size: u32) -> OverlapBlocks {
164- &self,
165- token_ids: &[i64],
166- block_size: u32,
167- cfg: &ScoringConfig,
168- ) -> OverlapScores {
169 let t_hash = std::time::Instant::now();215 let t_hash = std::time::Instant::now();
170 let block_hashes = compute_block_hash_for_seq(token_ids, block_size);216 let block_hashes = compute_block_hash_for_seq(token_ids, block_size);
171 let hash_us = t_hash.elapsed().as_micros();217 let hash_us = t_hash.elapsed().as_micros();
172 218 
173- let scores = self.find_matches_by_hash(&block_hashes, cfg);219+ let overlap = self.find_matches_by_hash(&block_hashes);
174 220 
175 tracing::debug!(221 tracing::debug!(
176 num_tokens = token_ids.len(),222 num_tokens = token_ids.len(),
177 block_size,223 block_size,
178 num_hashes = block_hashes.len(),224 num_hashes = block_hashes.len(),
179 hash_us,225 hash_us,
180- scores = scores.scores.len(),226+ matched = overlap.blocks.len(),
181 "hash_computed"227 "hash_computed"
182 );228 );
183- scores229+ overlap
184 }230 }
185 231 
186- pub fn find_matches_by_hash(232+ pub fn find_matches_by_hash(&self, block_hashes: &[LocalBlockHash]) -> OverlapBlocks {
187- &self,233+ self.find_matches_with_coverage(block_hashes).0
188- block_hashes: &[LocalBlockHash],
189- cfg: &ScoringConfig,
190- ) -> OverlapScores {
191- let mut scores = OverlapScores::default();
192- 
193- let hbm_scores = self.hbm_tree.find_matches(block_hashes);
194- for (worker, depth) in hbm_scores.scores {
195- scores.add_score(worker, depth * cfg.hbm_weight);
196- }
197- 
198- if block_hashes.len() > FLAT_PAR_THRESHOLD {
199- self.flat_lookup_parallel(block_hashes, cfg, &mut scores);
200- } else {
201- self.flat_lookup_sequential(block_hashes, cfg, &mut scores);
202- }
203- 
204- scores
205 }234 }
206 235 
207- /// Sequential CPU + Disk flat lookup (for small hash sets).236+ /// Query with per-DP absolute coverage end (in blocks).
208- fn flat_lookup_sequential(237+ ///
238+ /// `npu_blocks` / `cpu_blocks` / `disk_blocks` remain per-medium segment
239+ /// lengths. `coverage_end` is the farthest absolute prefix end across
240+ /// media for each `(instance_id, dp_rank)` — used for `matched_tokens`
241+ /// so overlapping replicas (same prefix on NPU+CPU+Disk) are not summed.
242+ fn find_matches_with_coverage(
209 &self,243 &self,
210 block_hashes: &[LocalBlockHash],244 block_hashes: &[LocalBlockHash],
211- cfg: &ScoringConfig,245+ ) -> (OverlapBlocks, FxHashMap<(String, DpRank), u32>) {
212- scores: &mut OverlapScores,246+ let mut overlap = OverlapBlocks::default();
213- ) {247+ let mut coverage_end: FxHashMap<(String, DpRank), u32> = FxHashMap::default();
214- let cpu = self.cpu_blocks.read();248+ 
215- for hash in block_hashes {249+ // 1) HBM prefix match.
216- if let Some(workers) = cpu.get(hash) {250+ let hbm: FxHashMap<WorkerKey, PrefixMatch> =
217- for w in workers {251+ self.hbm_tree.find_matches_detailed(block_hashes);
218- scores.add_score(w.clone(), cfg.cpu_weight);252+ for (worker, m) in &hbm {
219- }253+ if m.depth == 0 {
254+ continue;
220 }255 }
256+ overlap.add_blocks(worker.clone(), m.depth);
257+ Self::note_coverage(
258+ &mut coverage_end,
259+ &worker.instance_id,
260+ worker.dp_rank,
261+ m.depth,
262+ );
221 }263 }
222- drop(cpu);
223 264 
224- let disk = self.disk_blocks.read();265+ // Breakpoints need last_seq_hash for continuation.
225- for hash in block_hashes {266+ let hbm_breaks: Vec<TierBreakpoint> = hbm
226- if let Some(workers) = disk.get(hash) {267+ .iter()
227- for w in workers {268+ .filter(|(_, m)| m.depth > 0)
228- scores.add_score(w.clone(), cfg.disk_weight);269+ .filter_map(|(w, m)| {
229- }270+ Some(TierBreakpoint {
230- }271+ instance_id: w.instance_id.clone(),
231- }272+ dp_rank: w.dp_rank,
232- }273+ end_pos: m.depth as usize,
274+ last_seq: m.last_seq_hash?,
275+ })
276+ })
277+ .collect();
233 278 
234- /// Parallel CPU + Disk flat lookup using rayon (for large hash sets,279+ // 2) CPU: continue from HBM breakpoints; root walk runs for every
235- /// e.g. DeepSeek V4 with 32K+ block hashes).280+ // worker owning the first edge (replicas of an already-covered
236- fn flat_lookup_parallel(281+ // prefix are reported as real segment lengths — coverage is a
237- &self,282+ // max, so they cannot inflate matched_tokens).
238- block_hashes: &[LocalBlockHash],283+ let cpu_hits = self.lower_tier_lookup(
239- cfg: &ScoringConfig,284+ block_hashes,
240- scores: &mut OverlapScores,285+ &hbm_breaks,
241- ) {286+ &self.cpu_tiers,
242- use rayon::prelude::*;287+ &mut overlap,
243- 288+ &mut coverage_end,
244- let cpu = self.cpu_blocks.read();
245- let disk = self.disk_blocks.read();
246- 
247- let (cpu_scores, disk_scores): (OverlapScores, OverlapScores) = rayon::join(
248- || {
249- block_hashes
250- .par_iter()
251- .fold(OverlapScores::default, |mut acc, hash| {
252- if let Some(workers) = cpu.get(hash) {
253- for w in workers {
254- acc.add_score(w.clone(), cfg.cpu_weight);
255- }
256- }
257- acc
258- })
259- .reduce(OverlapScores::default, |mut a, b| {
260- a.merge(b);
261- a
262- })
263- },
264- || {
265- block_hashes
266- .par_iter()
267- .fold(OverlapScores::default, |mut acc, hash| {
268- if let Some(workers) = disk.get(hash) {
269- for w in workers {
270- acc.add_score(w.clone(), cfg.disk_weight);
271- }
272- }
273- acc
274- })
275- .reduce(OverlapScores::default, |mut a, b| {
276- a.merge(b);
277- a
278- })
279- },
280 );289 );
281 290 
282- scores.merge(cpu_scores);291+ let cpu_breaks: Vec<TierBreakpoint> = cpu_hits
283- scores.merge(disk_scores);292+ .iter()
293+ .filter(|(_, h)| h.count > 0)
294+ .filter_map(|(w, h)| {
295+ Some(TierBreakpoint {
296+ instance_id: w.instance_id.clone(),
297+ dp_rank: w.dp_rank,
298+ end_pos: h.end_pos(),
299+ last_seq: h.last_matched_hash?,
300+ })
301+ })
302+ .collect();
303+ 
304+ // 3) Disk: continue from max(HBM, CPU) per DP (CPU wins when it
305+ // extends further — matches vLLM lookup: CPU then Disk after NPU).
306+ let disk_breaks = Self::merge_tier_breakpoints(&hbm_breaks, &cpu_breaks);
307+ self.lower_tier_lookup(
308+ block_hashes,
309+ &disk_breaks,
310+ &self.disk_tiers,
311+ &mut overlap,
312+ &mut coverage_end,
313+ );
314+ 
315+ (overlap, coverage_end)
316+ }
317+ 
318+ #[inline]
319+ fn note_coverage(
320+ coverage_end: &mut FxHashMap<(String, DpRank), u32>,
321+ instance_id: &str,
322+ dp_rank: DpRank,
323+ end: u32,
324+ ) {
325+ coverage_end
326+ .entry((instance_id.to_string(), dp_rank))
327+ .and_modify(|e| *e = (*e).max(end))
328+ .or_insert(end);
329+ }
330+ 
331+ /// Per `(instance_id, dp_rank)`, keep the farther breakpoint.
332+ ///
333+ /// `preferred` (CPU) overwrites `fallback` (HBM) when ``end_pos`` is
334+ /// greater or equal — so Disk resumes after the longest upstream prefix.
335+ fn merge_tier_breakpoints(
336+ fallback: &[TierBreakpoint],
337+ preferred: &[TierBreakpoint],
338+ ) -> Vec<TierBreakpoint> {
339+ let mut best: HashMap<(String, DpRank), TierBreakpoint> = HashMap::new();
340+ for b in fallback {
341+ best.insert((b.instance_id.clone(), b.dp_rank), b.clone());
342+ }
343+ for b in preferred {
344+ let key = (b.instance_id.clone(), b.dp_rank);
345+ match best.get(&key) {
346+ Some(existing) if b.end_pos < existing.end_pos => {}
347+ _ => {
348+ best.insert(key, b.clone());
349+ }
350+ }
351+ }
352+ best.into_values().collect()
353+ }
354+ 
355+ /// Build continuations and count contiguous lower-tier hits.
356+ ///
357+ /// - Root walks run for **every** worker owning the first edge — replicas
358+ /// of a prefix already covered upstream are reported as real segment
359+ /// lengths. This cannot inflate `matched_tokens` (coverage is a max
360+ /// over absolute ends, each bounded by the query length).
361+ /// - Continuation starts from each upstream ``TierBreakpoint``; a worker
362+ /// may hold several candidates (root + breakpoints), and the one with
363+ /// the farthest absolute end wins inside
364+ /// [`LowerTierIndexer::query_contiguous_hits`].
365+ fn lower_tier_lookup(
366+ &self,
367+ block_hashes: &[LocalBlockHash],
368+ upstream_breaks: &[TierBreakpoint],
369+ tiers: &LowerTierIndexer,
370+ overlap: &mut OverlapBlocks,
371+ coverage_end: &mut FxHashMap<(String, DpRank), u32>,
372+ ) -> FxHashMap<WorkerKey, ContiguousHit> {
373+ if block_hashes.is_empty() {
374+ return FxHashMap::default();
375+ }
376+ 
377+ let mut continuations: FxHashMap<WorkerKey, Vec<LowerTierContinuation>> =
378+ FxHashMap::default();
379+ 
380+ // Root walk: unconditional, so a longer replica on this tier is never
381+ // hidden by an upstream (possibly shorter) hit.
382+ for w in tiers.root_workers(block_hashes[0]) {
383+ continuations
384+ .entry(w)
385+ .or_default()
386+ .push(LowerTierContinuation::from_root(0));
387+ }
388+ 
389+ // Continue from each upstream breakpoint (candidate list — the walk
390+ // keeps the farthest end per worker).
391+ for b in upstream_breaks {
392+ if b.end_pos >= block_hashes.len() {
393+ continue;
394+ }
395+ for w in tiers.edge_owners(Some(b.last_seq), block_hashes[b.end_pos]) {
396+ continuations
397+ .entry(w)
398+ .or_default()
399+ .push(LowerTierContinuation::new(b.end_pos, b.last_seq));
400+ }
401+ }
402+ 
403+ if continuations.is_empty() {
404+ return FxHashMap::default();
405+ }
406+ 
407+ let hits = tiers.query_contiguous_hits(block_hashes, &continuations);
408+ for (worker, hit) in &hits {
409+ if hit.count > 0 {
410+ overlap.add_blocks(worker.clone(), hit.count as u32);
411+ Self::note_coverage(
412+ coverage_end,
413+ &worker.instance_id,
414+ worker.dp_rank,
415+ hit.end_pos() as u32,
416+ );
417+ }
418+ }
419+ hits
284 }420 }
285 421 
286 // -----------------------------------------------------------------------422 // -----------------------------------------------------------------------
@@ -289,32 +425,54 @@ impl IndexerEntry {
289 425 
290 /// Ingest offload blocks from vLLM non-HBM events.426 /// Ingest offload blocks from vLLM non-HBM events.
291 ///427 ///
292- /// For each `(block_hash, tokens_hash)` pair, checks whether there are428+ /// Each triple is `(block_hash, tokens_hash, parent_hash)` where
293- /// pending pool backend events waiting for this block. Matched entries429+ /// `parent_hash` is the immediately preceding engine `block_hash` in
294- /// are removed from `pending_pool` and returned (grouped by worker) so430+ /// this offload chain (or the chain's original `parent_block_hash` for
295- /// the caller can insert them into the radix tree / flat store.431+ /// the first block). Checks whether there are pending pool backend
296- /// Unmatched entries are cached in `offload` for later pool confirmation.432+ /// events waiting for each block. Matched entries are removed from
433+ /// `pending_pool` and returned (grouped by worker). Content is retained
434+ /// (always retained; TTL-swept via [`CONTENT_TTL`]). Unmatched
435+ /// entries are cached in `offload`.
297 pub fn ingest_offload_blocks(436 pub fn ingest_offload_blocks(
298 &self,437 &self,
299- pairs: &[(u64, u64)],438+ triples: &[(u64, u64, Option<u64>)],
300- ) -> HashMap<WorkerKey, Vec<KvCacheStoredBlockData>> {439+ ) -> HashMap<WorkerKey, Vec<(Option<u64>, KvCacheStoredBlockData)>> {
301 let matched = {440 let matched = {
302 let mut state = self.offload_pool_state.write();441 let mut state = self.offload_pool_state.write();
303- let mut matched: HashMap<WorkerKey, Vec<KvCacheStoredBlockData>> = HashMap::new();442+ let mut matched: HashMap<WorkerKey, Vec<(Option<u64>, KvCacheStoredBlockData)>> =
443+ HashMap::new();
304 444 
305- for &(block_hash, tokens_hash) in pairs {445+ for &(block_hash, tokens_hash, parent_hash) in triples {
446+ let cache_entry = OffloadCacheEntry {
447+ tokens_hash,
448+ parent_hash,
449+ };
306 if let Some(pending) = state.pending_pool.remove(&block_hash) {450 if let Some(pending) = state.pending_pool.remove(&block_hash) {
307- for entry in pending {451+ state.content.insert(
308- matched452+ block_hash,
309- .entry(entry.worker)453+ ContentEntry {
310- .or_default()454+ entry: cache_entry,
311- .push(KvCacheStoredBlockData {455+ inserted_at: std::time::Instant::now(),
456+ },
457+ );
458+ state.offload.remove(&block_hash);
459+ for pending_entry in pending {
460+ matched.entry(pending_entry.worker).or_default().push((
461+ parent_hash,
462+ KvCacheStoredBlockData {
312 block_hash,463 block_hash,
313 tokens_hash,464 tokens_hash,
314- });465+ },
466+ ));
315 }467 }
316 } else {468 } else {
317- state.offload.insert(block_hash, tokens_hash);469+ // Unconfirmed offload: no TTL (pool confirmation may be
470+ // arbitrarily late). If this hash was previously confirmed
471+ // and re-offloaded, drop the stale content mapping to keep
472+ // the offload/content invariant — the tier still carries
473+ // the mapping for later pool events.
474+ state.content.remove(&block_hash);
475+ state.offload.insert(block_hash, cache_entry);
318 }476 }
319 }477 }
320 matched478 matched
@@ -323,73 +481,138 @@ impl IndexerEntry {
323 matched481 matched
324 }482 }
325 483 
326- /// Ingest pool backend blocks from Mooncake stored events.484+ /// Resolve `(parent_hash, tokens_hash)` from offload / retained content.
327 ///485 ///
328- /// For each `block_hash`, checks the `offload` cache for a pre-computed486+ /// When taking from `offload`, the mapping is moved into `content` (with a
329- /// `tokens_hash`. Matched entries are removed from `offload` and returned487+ /// fresh [`CONTENT_TTL`] window) so a later pool medium can reuse it. Tier
330- /// as `KvCacheStoredBlockData` blocks for radix tree insertion.488+ /// lookups happen outside this lock in [`Self::ingest_pool_blocks`].
331- /// Unmatched entries are queued in `pending_pool` so that a future489+ fn resolve_pool_content(
332- /// offload event can complete the match.490+ &self,
491+ state: &mut OffloadPoolState,
492+ block_hash: u64,
493+ ) -> Option<(Option<u64>, u64)> {
494+ if let Some(cached) = state.offload.remove(&block_hash) {
495+ state.content.insert(
496+ block_hash,
497+ ContentEntry {
498+ entry: cached,
499+ inserted_at: std::time::Instant::now(),
500+ },
501+ );
502+ return Some((cached.parent_hash, cached.tokens_hash));
503+ }
504+ if let Some(cached) = state.content.get(&block_hash) {
505+ return Some((cached.entry.parent_hash, cached.entry.tokens_hash));
506+ }
507+ None
508+ }
509+ 
510+ /// Ingest pool backend blocks from Mooncake / YuanRong stored events.
511+ ///
512+ /// For each `block_hash`, resolves `(tokens_hash, parent_hash)` from the
513+ /// offload cache, retained content map, or an already-indexed lower tier.
514+ /// Unmatched entries are queued in `pending_pool`.
333 pub fn ingest_pool_blocks(515 pub fn ingest_pool_blocks(
334 &self,516 &self,
335 block_hashes: &[u64],517 block_hashes: &[u64],
336 worker: &WorkerKey,518 worker: &WorkerKey,
337- ) -> Vec<KvCacheStoredBlockData> {519+ ) -> Vec<(Option<u64>, KvCacheStoredBlockData)> {
338- let matched = {520+ let mut matched = Vec::with_capacity(block_hashes.len());
339- let mut state = self.offload_pool_state.write();521+ let mut need_tier_lookup = Vec::new();
340- let mut matched = Vec::with_capacity(block_hashes.len());
341 522 
523+ {
524+ let mut state = self.offload_pool_state.write();
342 for &bh in block_hashes {525 for &bh in block_hashes {
343- if let Some(tokens_hash) = state.offload.remove(&bh) {526+ if let Some((parent_hash, tokens_hash)) = self.resolve_pool_content(&mut state, bh)
344- matched.push(KvCacheStoredBlockData {527+ {
345- block_hash: bh,528+ matched.push((
346- tokens_hash,529+ parent_hash,
347- });530+ KvCacheStoredBlockData {
531+ block_hash: bh,
532+ tokens_hash,
533+ },
534+ ));
348 } else {535 } else {
349- state536+ need_tier_lookup.push(bh);
350- .pending_pool
351- .entry(bh)
352- .or_default()
353- .insert(PendingPoolEvent {
354- worker: worker.clone(),
355- inserted_at: std::time::Instant::now(),
356- });
357 }537 }
358 }538 }
359- matched539+ } // release offload_pool_state before touching tier locks
360- }; // lock released before sweep540+ 
541+ for &bh in &need_tier_lookup {
542+ if let Some((parent, tokens)) = self
543+ .cpu_tiers
544+ .lookup_block(bh)
545+ .or_else(|| self.disk_tiers.lookup_block(bh))
546+ {
547+ let mut state = self.offload_pool_state.write();
548+ state.content.insert(
549+ bh,
550+ ContentEntry {
551+ entry: OffloadCacheEntry {
552+ tokens_hash: tokens,
553+ parent_hash: parent,
554+ },
555+ inserted_at: std::time::Instant::now(),
556+ },
557+ );
558+ matched.push((
559+ parent,
560+ KvCacheStoredBlockData {
561+ block_hash: bh,
562+ tokens_hash: tokens,
563+ },
564+ ));
565+ } else {
566+ let mut state = self.offload_pool_state.write();
567+ state
568+ .pending_pool
569+ .entry(bh)
570+ .or_default()
571+ .insert(PendingPoolEvent {
572+ worker: worker.clone(),
573+ inserted_at: std::time::Instant::now(),
574+ });
575+ }
576+ }
577+ 
361 self.maybe_sweep_pending();578 self.maybe_sweep_pending();
362 matched579 matched
363 }580 }
364 581 
365- /// Evict blocks from both caches (for removal events).582+ /// Evict blocks from pending caches (for removal events).
366 ///583 ///
367- /// Returns `block_hashes` that were in **neither** cache — these were584+ /// Returns `block_hashes` that need lower-tier / tree removal — already
368- /// already matched and inserted into the tree, so the caller must apply585+ /// confirmed into a tier, or with no matching pending state for this
369- /// tree removal for them.586+ /// worker (tier removal is a no-op if the hash never entered a tier).
370- ///587+ /// Unconfirmed `offload` / `pending_pool`-only entries are dropped
371- /// Blocks found in either cache were still pending and are simply evicted588+ /// without tier removal. Retained `content` is left for the migration
372- /// (they never entered the tree). This also cleans up `pending_pool`589+ /// window (bounded by [`CONTENT_TTL`]) and may resolve a later Disk store.
373- /// entries for the specific `worker`.
374 pub fn evict_pending_blocks(&self, block_hashes: &[u64], worker: &WorkerKey) -> Vec<u64> {590 pub fn evict_pending_blocks(&self, block_hashes: &[u64], worker: &WorkerKey) -> Vec<u64> {
375 let mut state = self.offload_pool_state.write();591 let mut state = self.offload_pool_state.write();
376 let mut need_tree_removal = Vec::new();592 let mut need_tree_removal = Vec::new();
377 593 
378 for &bh in block_hashes {594 for &bh in block_hashes {
379- // Try offload cache first595+ // Unconfirmed offload — never entered a tier.
380 if state.offload.remove(&bh).is_some() {596 if state.offload.remove(&bh).is_some() {
381- // Was waiting for pool event — now cancelled597+ state.content.remove(&bh);
382 continue;598 continue;
383 }599 }
384- // Try pending pool events600+ // Pool-first pending for this worker — may or may not be in a tier.
385 if let Some(entries) = state.pending_pool.get_mut(&bh) {601 if let Some(entries) = state.pending_pool.get_mut(&bh) {
602+ let before = entries.len();
386 entries.retain(|e| e.worker != *worker);603 entries.retain(|e| e.worker != *worker);
387- if entries.is_empty() {604+ if entries.len() != before {
388- state.pending_pool.remove(&bh);605+ if entries.is_empty() {
606+ state.pending_pool.remove(&bh);
607+ }
608+ // Content means another medium already confirmed this hash.
609+ if state.content.contains_key(&bh) {
610+ need_tree_removal.push(bh);
611+ }
612+ continue;
389 }613 }
390- continue;
391 }614 }
392- // Not in either cache — must already be in the tree615+ // Already confirmed (or unknown) — remove from the tier.
393 need_tree_removal.push(bh);616 need_tree_removal.push(bh);
394 }617 }
395 need_tree_removal618 need_tree_removal
@@ -398,11 +621,14 @@ impl IndexerEntry {
398 /// Remove all pending entries for a worker from `pending_pool`.621 /// Remove all pending entries for a worker from `pending_pool`.
399 ///622 ///
400 /// Called on worker disconnect / Cleared events. Returns the number of623 /// Called on worker disconnect / Cleared events. Returns the number of
401- /// block hashes that were affected.624+ /// block hashes whose `pending_pool` entries were fully cleared.
402 ///625 ///
403- /// Note: the `offload` map has no per-worker association, so offload626+ /// Note: `offload` / `content` have no per-worker association. Content is
404- /// entries are not purged here — they are bounded by the number of627+ /// TTL-evicted via [`Self::sweep_stale_pending`] ([`CONTENT_TTL`]); it is
405- /// outstanding offloaded blocks and evicted on match or TTL sweep.628+ /// kept across tier clears so a later Disk store can still promote the
629+ /// block. Unconfirmed `offload` entries are not TTL'd — they are bounded
630+ /// by the number of outstanding offloaded blocks and cleared on match,
631+ /// removal, or the pool confirmation eventually arriving.
406 pub fn remove_pending_worker(&self, worker: &WorkerKey) -> usize {632 pub fn remove_pending_worker(&self, worker: &WorkerKey) -> usize {
407 let mut state = self.offload_pool_state.write();633 let mut state = self.offload_pool_state.write();
408 let mut removed = 0usize;634 let mut removed = 0usize;
@@ -419,27 +645,38 @@ impl IndexerEntry {
419 removed645 removed
420 }646 }
421 647 
422- /// Total number of pending entries across both caches (for diagnostics).648+ /// Total pending + retained content entries (for diagnostics).
423 pub fn pending_count(&self) -> usize {649 pub fn pending_count(&self) -> usize {
424 let state = self.offload_pool_state.read();650 let state = self.offload_pool_state.read();
425- state.offload.len() + state.pending_pool.len()651+ state.offload.len() + state.pending_pool.len() + state.content.len()
426 }652 }
427 653 
428- /// Sweep stale `pending_pool` entries that exceed `ttl`, returning the654+ /// Sweep stale `pending_pool` and retained `content` entries that exceed
429- /// number evicted.655+ /// their TTLs, returning the total number of entries evicted.
430 ///656 ///
431- /// Only `pending_pool` entries are swept (they carry per-entry timestamps).657+ /// Unconfirmed `offload` entries are **not** swept: pool confirmation may
432- /// `offload` entries have no timestamp and are kept until matched by a658+ /// arrive arbitrarily late, so they are bounded only by the number of
433- /// pool event or explicitly evicted by a removal event — they are bounded659+ /// outstanding offloaded blocks (original two-phase design). `content`
434- /// by the number of outstanding offloaded blocks.660+ /// must survive lower-tier removal for the CPU→Disk migration window,
435- pub fn sweep_stale_pending(&self, ttl: std::time::Duration) -> usize {661+ /// so it is bounded by [`CONTENT_TTL`] instead.
662+ pub fn sweep_stale_pending(
663+ &self,
664+ pending_ttl: std::time::Duration,
665+ content_ttl: std::time::Duration,
666+ ) -> usize {
436 let mut state = self.offload_pool_state.write();667 let mut state = self.offload_pool_state.write();
437 let mut pruned = 0usize;668 let mut pruned = 0usize;
438 let now = std::time::Instant::now();669 let now = std::time::Instant::now();
439 670 
671+ let before_content = state.content.len();
672+ state
673+ .content
674+ .retain(|_bh, e| now.duration_since(e.inserted_at) < content_ttl);
675+ pruned += before_content - state.content.len();
676+ 
440 state.pending_pool.retain(|_bh, entries| {677 state.pending_pool.retain(|_bh, entries| {
441 entries.retain(|e| {678 entries.retain(|e| {
442- let keep = now.duration_since(e.inserted_at) < ttl;679+ let keep = now.duration_since(e.inserted_at) < pending_ttl;
443 if !keep {680 if !keep {
444 pruned += 1;681 pruned += 1;
445 }682 }
@@ -452,8 +689,9 @@ impl IndexerEntry {
452 tracing::debug!(689 tracing::debug!(
453 pruned,690 pruned,
454 remaining_offload = state.offload.len(),691 remaining_offload = state.offload.len(),
692+ remaining_content = state.content.len(),
455 remaining_pending_keys = state.pending_pool.len(),693 remaining_pending_keys = state.pending_pool.len(),
456- "swept stale pending pool entries"694+ "swept stale content/pending pool entries"
457 );695 );
458 }696 }
459 697 
@@ -461,109 +699,19 @@ impl IndexerEntry {
461 }699 }
462 700 
463 // -----------------------------------------------------------------------701 // -----------------------------------------------------------------------
464- // Flat (CPU/Disk) store / remove helpers702+ // Event application
465 // -----------------------------------------------------------------------703 // -----------------------------------------------------------------------
466 704 
467- /// Insert a block into the CPU or Disk flat store.705+ /// Periodically sweep stale pending pool / content entries on every
468- fn flat_store(706+ /// `PENDING_SWEEP_THRESHOLD`-th ingest. Called after each `ingest_*`
469- blocks: &RwLock<FxHashMap<LocalBlockHash, FxHashSet<WorkerKey>>>,707+ /// operation.
470- lookups: &RwLock<FxHashMap<WorkerKey, FlatLookup>>,
471- worker: &WorkerKey,
472- tokens_hash: u64,
473- seq_hash: u64,
474- ) {
475- // Update flat block set: tokens_hash → {workers}
476- {
477- let mut map = blocks.write();
478- map.entry(LocalBlockHash(tokens_hash))
479- .or_default()
480- .insert(worker.clone());
481- }
482- // Update per-worker reverse lookup for removal
483- {
484- let mut lu = lookups.write();
485- lu.entry(worker.clone())
486- .or_default()
487- .insert(seq_hash, tokens_hash);
488- }
489- tracing::trace!(
490- instance_id = %worker.instance_id,
491- dp_rank = worker.dp_rank,
492- medium = %worker.medium.as_str(),
493- ?seq_hash,
494- tokens_hash,
495- "flat store"
496- );
497- }
498- 
499- /// Remove a block from the CPU or Disk flat store.
500- fn flat_remove(
501- blocks: &RwLock<FxHashMap<LocalBlockHash, FxHashSet<WorkerKey>>>,
502- lookups: &RwLock<FxHashMap<WorkerKey, FlatLookup>>,
503- worker: &WorkerKey,
504- seq_hash: u64,
505- ) {
506- // Find tokens_hash from per-worker reverse lookup
507- let tokens_hash = {
508- let lu = lookups.read();
509- lu.get(worker).and_then(|m| m.get(&seq_hash).copied())
510- };
511- 
512- if let Some(th) = tokens_hash {
513- // Remove worker from the flat set
514- let mut map = blocks.write();
515- if let Some(set) = map.get_mut(&LocalBlockHash(th)) {
516- set.remove(worker);
517- if set.is_empty() {
518- map.remove(&LocalBlockHash(th));
519- }
520- }
521- // Clean up reverse lookup
522- let mut lu = lookups.write();
523- if let Some(m) = lu.get_mut(worker) {
524- m.remove(&seq_hash);
525- }
526- }
527- }
528- 
529- /// Clear all CPU/Disk flat blocks for a worker.
530- fn flat_clear(
531- blocks: &RwLock<FxHashMap<LocalBlockHash, FxHashSet<WorkerKey>>>,
532- lookups: &RwLock<FxHashMap<WorkerKey, FlatLookup>>,
533- worker: &WorkerKey,
534- ) {
535- // Collect tokens_hashes from reverse lookup
536- let tokens_hashes: Vec<LocalBlockHash> = {
537- let lu = lookups.read();
538- lu.get(worker)
539- .map(|m| m.values().map(|&th| LocalBlockHash(th)).collect())
540- .unwrap_or_default()
541- };
542- // Remove worker from each block set
543- if !tokens_hashes.is_empty() {
544- let mut map = blocks.write();
545- for th in &tokens_hashes {
546- if let Some(set) = map.get_mut(th) {
547- set.remove(worker);
548- if set.is_empty() {
549- map.remove(th);
550- }
551- }
552- }
553- }
554- // Clear reverse lookup
555- lookups.write().remove(worker);
556- }
557- 
558- /// Periodically sweep stale pending pool entries when the ingest counter
559- /// exceeds the threshold. Called after each `ingest_*` operation.
560 fn maybe_sweep_pending(&self) {708 fn maybe_sweep_pending(&self) {
561 let count = self709 let count = self
562 .pending_ingest_count710 .pending_ingest_count
563 .fetch_add(1, std::sync::atomic::Ordering::Relaxed)711 .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
564 + 1;712 + 1;
565 if count.is_multiple_of(PENDING_SWEEP_THRESHOLD) {713 if count.is_multiple_of(PENDING_SWEEP_THRESHOLD) {
566- let pruned = self.sweep_stale_pending(PENDING_TTL);714+ let pruned = self.sweep_stale_pending(PENDING_TTL, CONTENT_TTL);
567 if pruned > 0 {715 if pruned > 0 {
568 tracing::debug!(716 tracing::debug!(
569 pruned,717 pruned,
@@ -596,74 +744,58 @@ impl IndexerEntry {
596 match event {744 match event {
597 KvCacheEventData::Stored(store_data) => {745 KvCacheEventData::Stored(store_data) => {
598 match worker.medium {746 match worker.medium {
599- StorageMedium::Xpu | StorageMedium::Unknown => {747+ StorageMedium::Npu | StorageMedium::Unknown => {
600 // HBM: prefix-chain tree insert748 // HBM: prefix-chain tree insert
601 let mut lookups = self.lookups.write();749 let mut lookups = self.lookups.write();
602 let lookup = lookups.entry(worker.clone()).or_default();750 let lookup = lookups.entry(worker.clone()).or_default();
603 self.hbm_tree.apply_store(worker, lookup, store_data)751 self.hbm_tree.apply_store(worker, lookup, store_data)
604 }752 }
605 StorageMedium::Cpu => {753 StorageMedium::Cpu => {
606- // CPU: flat insert754+ self.cpu_tiers.store_blocks(worker, store_data);
607- for block in &store_data.blocks {
608- Self::flat_store(
609- &self.cpu_blocks,
610- &self.cpu_lookups,
611- worker,
612- block.tokens_hash,
613- block.block_hash,
614- );
615- }
616 Ok(())755 Ok(())
617 }756 }
618 StorageMedium::Disk => {757 StorageMedium::Disk => {
619- // Disk: flat insert758+ self.disk_tiers.store_blocks(worker, store_data);
620- for block in &store_data.blocks {
621- Self::flat_store(
622- &self.disk_blocks,
623- &self.disk_lookups,
624- worker,
625- block.tokens_hash,
626- block.block_hash,
627- );
628- }
629 Ok(())759 Ok(())
630 }760 }
631 }761 }
632 }762 }
633 KvCacheEventData::Removed { block_hashes } => match worker.medium {763 KvCacheEventData::Removed { block_hashes } => match worker.medium {
634- StorageMedium::Xpu | StorageMedium::Unknown => {764+ StorageMedium::Npu | StorageMedium::Unknown => {
635 let mut lookups = self.lookups.write();765 let mut lookups = self.lookups.write();
636 let lookup = lookups.entry(worker.clone()).or_default();766 let lookup = lookups.entry(worker.clone()).or_default();
637 let result = self.hbm_tree.apply_remove(worker, lookup, block_hashes);767 let result = self.hbm_tree.apply_remove(worker, lookup, block_hashes);
638 self.maybe_sweep_hbm();768 self.maybe_sweep_hbm();
639 result769 result
640 }770 }
771+ // Retained content is deliberately NOT pruned here: a CPU
772+ // eviction may be the pool migrating the block to Disk, and
773+ // the Disk store event must still resolve via `content`
774+ // (bounded by `CONTENT_TTL` sweep instead).
641 StorageMedium::Cpu => {775 StorageMedium::Cpu => {
642- for &h in block_hashes {776+ self.cpu_tiers.remove_blocks(worker, block_hashes);
643- Self::flat_remove(&self.cpu_blocks, &self.cpu_lookups, worker, h);
644- }
645 Ok(())777 Ok(())
646 }778 }
647 StorageMedium::Disk => {779 StorageMedium::Disk => {
648- for &h in block_hashes {780+ self.disk_tiers.remove_blocks(worker, block_hashes);
649- Self::flat_remove(&self.disk_blocks, &self.disk_lookups, worker, h);
650- }
651 Ok(())781 Ok(())
652 }782 }
653 },783 },
654 KvCacheEventData::Cleared => {784 KvCacheEventData::Cleared => {
655 match worker.medium {785 match worker.medium {
656- StorageMedium::Xpu | StorageMedium::Unknown => {786+ StorageMedium::Npu | StorageMedium::Unknown => {
657 let mut lookups = self.lookups.write();787 let mut lookups = self.lookups.write();
658 let lookup = lookups.entry(worker.clone()).or_default();788 let lookup = lookups.entry(worker.clone()).or_default();
659 self.hbm_tree.remove_worker(worker, lookup);789 self.hbm_tree.remove_worker(worker, lookup);
660 self.maybe_sweep_hbm();790 self.maybe_sweep_hbm();
661 }791 }
792+ // Same as Removed: retained content survives the clear so
793+ // a later Disk store can still promote these blocks.
662 StorageMedium::Cpu => {794 StorageMedium::Cpu => {
663- Self::flat_clear(&self.cpu_blocks, &self.cpu_lookups, worker);795+ self.cpu_tiers.clear_worker(worker);
664 }796 }
665 StorageMedium::Disk => {797 StorageMedium::Disk => {
666- Self::flat_clear(&self.disk_blocks, &self.disk_lookups, worker);798+ self.disk_tiers.clear_worker(worker);
667 }799 }
668 }800 }
669 Ok(())801 Ok(())
@@ -689,39 +821,21 @@ impl IndexerEntry {
689 lookups.remove(wk);821 lookups.remove(wk);
690 }822 }
691 }823 }
692- // CPU flat — collect keys first to avoid holding read lock824+ // CPU / Disk continuation-edge indexes
693- // across the write in flat_clear.825+ for wk in self.cpu_tiers.worker_keys() {
694- {826+ if wk.instance_id == instance_id && wk.dp_rank == dp_rank {
695- let cpu_matches: Vec<WorkerKey> = {827+ self.cpu_tiers.clear_worker(&wk);
696- let cpu_lu = self.cpu_lookups.read();
697- cpu_lu
698- .keys()
699- .filter(|wk| wk.instance_id == instance_id && wk.dp_rank == dp_rank)
700- .cloned()
701- .collect()
702- };
703- for wk in &cpu_matches {
704- Self::flat_clear(&self.cpu_blocks, &self.cpu_lookups, wk);
705 }828 }
706 }829 }
707- // Disk flat — same pattern830+ for wk in self.disk_tiers.worker_keys() {
708- {831+ if wk.instance_id == instance_id && wk.dp_rank == dp_rank {
709- let disk_matches: Vec<WorkerKey> = {832+ self.disk_tiers.clear_worker(&wk);
710- let disk_lu = self.disk_lookups.read();
711- disk_lu
712- .keys()
713- .filter(|wk| wk.instance_id == instance_id && wk.dp_rank == dp_rank)
714- .cloned()
715- .collect()
716- };
717- for wk in &disk_matches {
718- Self::flat_clear(&self.disk_blocks, &self.disk_lookups, wk);
719 }833 }
720 }834 }
721 // Offload/pool pending state — clean up pool entries waiting for835 // Offload/pool pending state — clean up pool entries waiting for
722- // this worker. Offload entries have no per-worker association so836+ // this worker. Unconfirmed offload / retained content entries are
723- // we cannot selectively purge them here; they are bounded by the837+ // not per-worker: offload is bounded by outstanding offloaded blocks,
724- // number of outstanding offloaded blocks.838+ // and content is TTL-evicted (`CONTENT_TTL`), so both self-clean.
725 {839 {
726 let mut state = self.offload_pool_state.write();840 let mut state = self.offload_pool_state.write();
727 state.pending_pool.retain(|_, entries| {841 state.pending_pool.retain(|_, entries| {
@@ -735,26 +849,14 @@ impl IndexerEntry {
735 /// Get the total number of cached blocks across all workers and media.849 /// Get the total number of cached blocks across all workers and media.
736 pub fn total_blocks(&self) -> usize {850 pub fn total_blocks(&self) -> usize {
737 let hbm = self.lookups.read().values().map(|l| l.len()).sum::<usize>();851 let hbm = self.lookups.read().values().map(|l| l.len()).sum::<usize>();
738- let cpu = self852+ hbm + self.cpu_tiers.total_blocks() + self.disk_tiers.total_blocks()
739- .cpu_lookups
740- .read()
741- .values()
742- .map(|l| l.len())
743- .sum::<usize>();
744- let disk = self
745- .disk_lookups
746- .read()
747- .values()
748- .map(|l| l.len())
749- .sum::<usize>();
750- hbm + cpu + disk
751 }853 }
752 854 
753 /// Get all registered worker keys.855 /// Get all registered worker keys.
754 pub fn worker_keys(&self) -> Vec<WorkerKey> {856 pub fn worker_keys(&self) -> Vec<WorkerKey> {
755 let mut keys: Vec<WorkerKey> = self.lookups.read().keys().cloned().collect();857 let mut keys: Vec<WorkerKey> = self.lookups.read().keys().cloned().collect();
756- keys.extend(self.cpu_lookups.read().keys().cloned());858+ keys.extend(self.cpu_tiers.worker_keys());
757- keys.extend(self.disk_lookups.read().keys().cloned());859+ keys.extend(self.disk_tiers.worker_keys());
758 keys860 keys
759 }861 }
760}862}
@@ -762,14 +864,13 @@ impl IndexerEntry {
762/// Top-level indexer managing multiple (model, tenant) trees.864/// Top-level indexer managing multiple (model, tenant) trees.
763pub struct Indexer {865pub struct Indexer {
764 entries: DashMap<IndexerKey, Arc<IndexerEntry>>,866 entries: DashMap<IndexerKey, Arc<IndexerEntry>>,
765- scoring: ScoringConfig,
766}867}
767 868 
768impl Indexer {869impl Indexer {
769- pub fn new(scoring: ScoringConfig) -> Self {870+ /// Create an indexer.
871+ pub fn new() -> Self {
770 Self {872 Self {
771 entries: DashMap::new(),873 entries: DashMap::new(),
772- scoring,
773 }874 }
774 }875 }
775 876 
@@ -795,7 +896,8 @@ impl Indexer {
795 self.entries.get(&key).map(|e| e.value().clone())896 self.entries.get(&key).map(|e| e.value().clone())
796 }897 }
797 898 
798- /// Remove an indexer entry if it has no more workers across any medium.899+ /// Remove an indexer entry if it has no more workers across any medium
900+ /// and no pending offload / content / pool state remains.
799 pub fn remove_if_empty(&self, model_name: &str, tenant_id: &str) {901 pub fn remove_if_empty(&self, model_name: &str, tenant_id: &str) {
800 let key = IndexerKey {902 let key = IndexerKey {
801 model_name: model_name.to_string(),903 model_name: model_name.to_string(),
@@ -804,8 +906,8 @@ impl Indexer {
804 let should_remove = self.entries.get(&key).is_some_and(|e| {906 let should_remove = self.entries.get(&key).is_some_and(|e| {
805 let entry = e.value();907 let entry = e.value();
806 entry.lookups.read().is_empty()908 entry.lookups.read().is_empty()
807- && entry.cpu_lookups.read().is_empty()909+ && entry.cpu_tiers.is_empty()
808- && entry.disk_lookups.read().is_empty()910+ && entry.disk_tiers.is_empty()
809 && entry.pending_count() == 0911 && entry.pending_count() == 0
810 });912 });
811 if should_remove {913 if should_remove {
@@ -813,7 +915,7 @@ impl Indexer {
813 }915 }
814 }916 }
815 917 
816- /// Query overlap scores for a token sequence against a specific model/tenant.918+ /// Query matched block counts for a token sequence against a specific model/tenant.
817 ///919 ///
818 /// `block_size` determines the token-to-hash granularity — it must match920 /// `block_size` determines the token-to-hash granularity — it must match
819 /// the size used by the engine when publishing events.921 /// the size used by the engine when publishing events.
@@ -833,10 +935,21 @@ impl Indexer {
833 tenant_id: tenant_id.to_string(),935 tenant_id: tenant_id.to_string(),
834 })?;936 })?;
835 937 
836- let overlap = entry.find_matches(token_ids, block_size, &self.scoring);938+ let t_hash = std::time::Instant::now();
939+ let block_hashes = compute_block_hash_for_seq(token_ids, block_size);
940+ let hash_us = t_hash.elapsed().as_micros();
941+ let (overlap, coverage_end) = entry.find_matches_with_coverage(&block_hashes);
942+ tracing::debug!(
943+ num_tokens = token_ids.len(),
944+ block_size,
945+ num_hashes = block_hashes.len(),
946+ hash_us,
947+ matched = overlap.blocks.len(),
948+ "hash_computed"
949+ );
837 let t_tree = t0.elapsed();950 let t_tree = t0.elapsed();
838 951 
839- let resp = self.build_response(overlap, model_name, tenant_id, block_size);952+ let resp = Self::build_response(overlap, &coverage_end, model_name, tenant_id, block_size);
840 let total = t0.elapsed();953 let total = t0.elapsed();
841 954 
842 tracing::debug!(955 tracing::debug!(
@@ -849,7 +962,7 @@ impl Indexer {
849 resp962 resp
850 }963 }
851 964 
852- /// Query overlap scores using pre-computed `LocalBlockHash` values.965+ /// Query matched block counts using pre-computed `LocalBlockHash` values.
853 pub fn query_by_hash(966 pub fn query_by_hash(
854 &self,967 &self,
855 model_name: &str,968 model_name: &str,
@@ -863,16 +976,23 @@ impl Indexer {
863 tenant_id: tenant_id.to_string(),976 tenant_id: tenant_id.to_string(),
864 })?;977 })?;
865 978 
866- let overlap = entry.find_matches_by_hash(block_hashes, &self.scoring);979+ let (overlap, coverage_end) = entry.find_matches_with_coverage(block_hashes);
867 // Default to 1 token per hash (no scaling) since we don't know the980 // Default to 1 token per hash (no scaling) since we don't know the
868 // original block_size from the hash alone.981 // original block_size from the hash alone.
869- self.build_response(overlap, model_name, tenant_id, 1)982+ Self::build_response(overlap, &coverage_end, model_name, tenant_id, 1)
870 }983 }
871 984 
872- /// Build a `QueryResponse` from weighted overlap scores.985+ /// Build a `QueryResponse` from per-worker matched block counts.
986+ ///
987+ /// Per-medium fields are segment lengths. `matched_tokens` uses the
988+ /// farthest absolute coverage end across media for that DP (never the
989+ /// sum of segment lengths — every hit records a coverage end, so the
990+ /// defensive sum fallback below is unreachable), so NPU+CPU+Disk
991+ /// replicas of the same prefix do not inflate the cached prefix beyond
992+ /// the input.
873 fn build_response(993 fn build_response(
874- &self,994+ overlap: OverlapBlocks,
875- overlap: OverlapScores,995+ coverage_end: &FxHashMap<(String, DpRank), u32>,
876 model_name: &str,996 model_name: &str,
877 tenant_id: &str,997 tenant_id: &str,
878 block_size: u32,998 block_size: u32,
@@ -886,43 +1006,40 @@ impl Indexer {
886 1006 
887 let mut instance_data: HashMap<String, InstanceMatchData> = HashMap::new();1007 let mut instance_data: HashMap<String, InstanceMatchData> = HashMap::new();
888 1008 
889- for (worker, &score) in &overlap.scores {1009+ for (worker, &matched_blocks) in &overlap.blocks {
890 let dp_rank_str = worker.dp_rank.to_string();1010 let dp_rank_str = worker.dp_rank.to_string();
891- 
892- // Derive per-medium block count from the score and weight.
893- let matched_blocks = match worker.medium {
894- StorageMedium::Xpu | StorageMedium::Unknown => score / self.scoring.hbm_weight,
895- StorageMedium::Cpu => score / self.scoring.cpu_weight,
896- StorageMedium::Disk => score / self.scoring.disk_weight,
897- };
898- let matched_tokens = matched_blocks * block_size;
899- 
900 let imd = instance_data.entry(worker.instance_id.clone()).or_default();1011 let imd = instance_data.entry(worker.instance_id.clone()).or_default();
901- 1012+ let dp_match = imd.dp.entry(dp_rank_str).or_default();
902- imd.longest_matched = imd.longest_matched.max(matched_tokens);
903- 
904- let dp_score = imd.dp.entry(dp_rank_str).or_default();
905 match worker.medium {1013 match worker.medium {
906- StorageMedium::Xpu | StorageMedium::Unknown => {1014+ StorageMedium::Npu | StorageMedium::Unknown => {
907- dp_score.xpu_score = dp_score.xpu_score.max(score);1015+ dp_match.npu_blocks = dp_match.npu_blocks.max(matched_blocks);
908- dp_score.xpu_blocks = dp_score.xpu_blocks.max(matched_blocks);
909 }1016 }
910 StorageMedium::Cpu => {1017 StorageMedium::Cpu => {
911- dp_score.cpu_score = dp_score.cpu_score.max(score);1018+ dp_match.cpu_blocks = dp_match.cpu_blocks.max(matched_blocks);
912- dp_score.cpu_blocks = dp_score.cpu_blocks.max(matched_blocks);
913 }1019 }
914 StorageMedium::Disk => {1020 StorageMedium::Disk => {
915- dp_score.disk_score = dp_score.disk_score.max(score);1021+ dp_match.disk_blocks = dp_match.disk_blocks.max(matched_blocks);
916- dp_score.disk_blocks = dp_score.disk_blocks.max(matched_blocks);
917 }1022 }
918 }1023 }
919- dp_score.matched_tokens = dp_score.matched_tokens.max(matched_tokens);
920- dp_score.total = dp_score.xpu_score + dp_score.cpu_score + dp_score.disk_score;
921 }1024 }
922 1025 
923- // Compute total_score once per instance after all DPs have been populated.1026+ for (instance_id, imd) in instance_data.iter_mut() {
924- for imd in instance_data.values_mut() {1027+ let mut longest = 0u32;
925- imd.total_score = imd.dp.values().map(|s| s.total).sum();1028+ for (dp_rank_str, dp_match) in imd.dp.iter_mut() {
1029+ let dp_rank: DpRank = dp_rank_str.parse().unwrap_or(0);
1030+ let covered = coverage_end
1031+ .get(&(instance_id.clone(), dp_rank))
1032+ .copied()
1033+ .unwrap_or_else(|| {
1034+ dp_match
1035+ .npu_blocks
1036+ .saturating_add(dp_match.cpu_blocks)
1037+ .saturating_add(dp_match.disk_blocks)
1038+ });
1039+ dp_match.matched_tokens = covered.saturating_mul(block_size);
1040+ longest = longest.max(dp_match.matched_tokens);
1041+ }
1042+ imd.longest_matched = longest;
926 }1043 }
927 1044 
928 let mut response = QueryResponse::default();1045 let mut response = QueryResponse::default();
@@ -953,7 +1070,7 @@ impl Indexer {
953 1070 
954impl Default for Indexer {1071impl Default for Indexer {
955 fn default() -> Self {1072 fn default() -> Self {
956- Self::new(ScoringConfig::default())1073+ Self::new()
957 }1074 }
958}1075}
959 1076 
@@ -965,139 +1082,5 @@ pub struct IndexerSummary {
965 pub total_blocks: usize,1082 pub total_blocks: usize,
966}1083}
967 1084 
968-use serde::Serialize;
969- 
970#[cfg(test)]1085#[cfg(test)]
971-mod tests {1086+mod tests;
972- use super::*;
973- 
974- #[test]
975- fn test_indexer_get_or_create_and_query() {
976- let indexer = Indexer::new(ScoringConfig::default());
977- let entry = indexer.get_or_create("model-a", "tenant-1");
978- 
979- // Compute the actual hash for the test token sequence
980- let tokens: Vec<i64> = vec![10, 20, 30, 40];
981- let hashes = compute_block_hash_for_seq(&tokens, 4);
982- assert!(!hashes.is_empty());
983- let tokens_hash = hashes[0];
984- 
985- // Insert a worker with XPU blocks using the real hash
986- let wk_xpu = WorkerKey {
987- instance_id: "inst-1".into(),
988- backend_id: "inst-1".into(),
989- dp_rank: 0,
990- medium: StorageMedium::Xpu,
991- };
992- 
993- let store = KvCacheEventData::Stored(KvCacheStoreData {
994- parent_hash: None,
995- start_position: None,
996- blocks: vec![KvCacheStoredBlockData {
997- block_hash: 100,
998- tokens_hash: tokens_hash.0,
999- }],
1000- });
1001- entry.apply_event(&wk_xpu, &store).unwrap();
1002- 
1003- // Query with the same tokens
1004- let resp = indexer.query("model-a", "tenant-1", &tokens, 4).unwrap();
1005- let tenant = &resp.tenants["tenant-1"];
1006- let imd = &tenant["inst-1"];
1007- let dp0 = &imd.dp["0"];
1008- assert!(dp0.xpu_blocks > 0, "should have XPU match");
1009- assert_eq!(dp0.cpu_blocks, 0);
1010- assert_eq!(dp0.disk_blocks, 0);
1011- assert!(dp0.matched_tokens > 0);
1012- assert_eq!(imd.longest_matched, dp0.matched_tokens);
1013- }
1014- 
1015- #[test]
1016- fn test_per_tier_aggregation() {
1017- let indexer = Indexer::new(ScoringConfig::default());
1018- let entry = indexer.get_or_create("model-b", "t1");
1019- 
1020- // Two different token sequences → different block hashes
1021- let tokens_a: Vec<i64> = vec![10, 20, 30, 40];
1022- let tokens_b: Vec<i64> = vec![50, 60, 70, 80];
1023- let hash_a = compute_block_hash_for_seq(&tokens_a, 4)[0];
1024- let hash_b = compute_block_hash_for_seq(&tokens_b, 4)[0];
1025- 
1026- // Worker 1: XPU blocks
1027- let wk1 = WorkerKey {
1028- instance_id: "inst-1".into(),
1029- backend_id: "inst-1".into(),
1030- dp_rank: 0,
1031- medium: StorageMedium::Xpu,
1032- };
1033- entry
1034- .apply_event(
1035- &wk1,
1036- &KvCacheEventData::Stored(KvCacheStoreData {
1037- parent_hash: None,
1038- start_position: None,
1039- blocks: vec![KvCacheStoredBlockData {
1040- block_hash: 100,
1041- tokens_hash: hash_a.0,
1042- }],
1043- }),
1044- )
1045- .unwrap();
1046- 
1047- // Worker 2: CPU blocks (different instance, different tokens)
1048- let wk2 = WorkerKey {
1049- instance_id: "inst-2".into(),
1050- backend_id: "mooncake-1".into(),
1051- dp_rank: 0,
1052- medium: StorageMedium::Cpu,
1053- };
1054- entry
1055- .apply_event(
1056- &wk2,
1057- &KvCacheEventData::Stored(KvCacheStoreData {
1058- parent_hash: None,
1059- start_position: None,
1060- blocks: vec![KvCacheStoredBlockData {
1061- block_hash: 200,
1062- tokens_hash: hash_b.0,
1063- }],
1064- }),
1065- )
1066- .unwrap();
1067- 
1068- // Query with tokens_a — should match inst-1 (XPU) only
1069- let resp = indexer.query("model-b", "t1", &tokens_a, 4).unwrap();
1070- let tenant = &resp.tenants["t1"];
1071- 
1072- let imd1 = &tenant["inst-1"];
1073- let dp0 = &imd1.dp["0"];
1074- assert!(
1075- dp0.xpu_blocks > 0,
1076- "inst-1 should have XPU match for tokens_a"
1077- );
1078- assert_eq!(dp0.cpu_blocks, 0, "inst-1 should have no CPU match");
1079- 
1080- // Query with tokens_b — should match inst-2 (CPU) only
1081- let resp = indexer.query("model-b", "t1", &tokens_b, 4).unwrap();
1082- let tenant = &resp.tenants["t1"];
1083- 
1084- let imd2 = &tenant["inst-2"];
1085- let dp2 = &imd2.dp["0"];
1086- assert_eq!(dp2.xpu_blocks, 0, "inst-2 should have no XPU match");
1087- assert!(
1088- dp2.cpu_blocks > 0,
1089- "inst-2 should have CPU match for tokens_b"
1090- );
1091- }
1092- 
1093- #[test]
1094- fn test_no_indexer_error() {
1095- let indexer = Indexer::new(ScoringConfig::default());
1096- let err = indexer.query("no-such-model", "default", &[1, 2, 3, 4], 4);
1097- assert!(err.is_err());
1098- assert!(matches!(
1099- err.unwrap_err(),
1100- KvConductorError::NoIndexer { .. }
1101- ));
1102- }
1103-}
Amotor/kv_conductor/src/indexer/tests.rs+852-0
@@ -0,0 +1,852 @@
1+// Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.
2+// MindIE is licensed under Mulan PSL v2.
3+// You can use this software according to the terms and conditions of the Mulan PSL v2.
4+// You may obtain a copy of Mulan PSL v2 at:
5+// http://license.coscl.org.cn/MulanPSL2
6+// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
7+// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
8+// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
9+// See the Mulan PSL v2 for more details.
10+ 
11+use super::*;
12+ 
13+#[test]
14+fn test_indexer_get_or_create_and_query() {
15+ let indexer = Indexer::new();
16+ let entry = indexer.get_or_create("model-a", "tenant-1");
17+ 
18+ // Compute the actual hash for the test token sequence
19+ let tokens: Vec<i64> = vec![10, 20, 30, 40];
20+ let hashes = compute_block_hash_for_seq(&tokens, 4);
21+ assert!(!hashes.is_empty());
22+ let tokens_hash = hashes[0];
23+ 
24+ // Insert a worker with NPU blocks using the real hash
25+ let wk_npu = WorkerKey {
26+ instance_id: "inst-1".into(),
27+ backend_id: "inst-1".into(),
28+ dp_rank: 0,
29+ medium: StorageMedium::Npu,
30+ };
31+ 
32+ let store = KvCacheEventData::Stored(KvCacheStoreData {
33+ parent_hash: None,
34+ start_position: None,
35+ blocks: vec![KvCacheStoredBlockData {
36+ block_hash: 100,
37+ tokens_hash: tokens_hash.0,
38+ }],
39+ });
40+ entry.apply_event(&wk_npu, &store).unwrap();
41+ 
42+ // Query with the same tokens
43+ let resp = indexer.query("model-a", "tenant-1", &tokens, 4).unwrap();
44+ let tenant = &resp.tenants["tenant-1"];
45+ let imd = &tenant["inst-1"];
46+ let dp0 = &imd.dp["0"];
47+ assert!(dp0.npu_blocks > 0, "should have NPU match");
48+ assert_eq!(dp0.cpu_blocks, 0);
49+ assert_eq!(dp0.disk_blocks, 0);
50+ assert_eq!(
51+ dp0.matched_tokens,
52+ (dp0.npu_blocks + dp0.cpu_blocks + dp0.disk_blocks) * 4
53+ );
54+ assert_eq!(imd.longest_matched, dp0.matched_tokens);
55+}
56+ 
57+#[test]
58+fn test_per_tier_aggregation() {
59+ let indexer = Indexer::new();
60+ let entry = indexer.get_or_create("model-b", "t1");
61+ 
62+ // Two different token sequences → different block hashes
63+ let tokens_a: Vec<i64> = vec![10, 20, 30, 40];
64+ let tokens_b: Vec<i64> = vec![50, 60, 70, 80];
65+ let hash_a = compute_block_hash_for_seq(&tokens_a, 4)[0];
66+ let hash_b = compute_block_hash_for_seq(&tokens_b, 4)[0];
67+ 
68+ // Worker 1: NPU blocks
69+ let wk1 = WorkerKey {
70+ instance_id: "inst-1".into(),
71+ backend_id: "inst-1".into(),
72+ dp_rank: 0,
73+ medium: StorageMedium::Npu,
74+ };
75+ entry
76+ .apply_event(
77+ &wk1,
78+ &KvCacheEventData::Stored(KvCacheStoreData {
79+ parent_hash: None,
80+ start_position: None,
81+ blocks: vec![KvCacheStoredBlockData {
82+ block_hash: 100,
83+ tokens_hash: hash_a.0,
84+ }],
85+ }),
86+ )
87+ .unwrap();
88+ 
89+ // Worker 2: CPU blocks (different instance, different tokens)
90+ let wk2 = WorkerKey {
91+ instance_id: "inst-2".into(),
92+ backend_id: "mooncake-1".into(),
93+ dp_rank: 0,
94+ medium: StorageMedium::Cpu,
95+ };
96+ entry
97+ .apply_event(
98+ &wk2,
99+ &KvCacheEventData::Stored(KvCacheStoreData {
100+ parent_hash: None,
101+ start_position: None,
102+ blocks: vec![KvCacheStoredBlockData {
103+ block_hash: 200,
104+ tokens_hash: hash_b.0,
105+ }],
106+ }),
107+ )
108+ .unwrap();
109+ 
110+ // Query with tokens_a — should match inst-1 (NPU) only
111+ let resp = indexer.query("model-b", "t1", &tokens_a, 4).unwrap();
112+ let tenant = &resp.tenants["t1"];
113+ 
114+ let imd1 = &tenant["inst-1"];
115+ let dp0 = &imd1.dp["0"];
116+ assert!(
117+ dp0.npu_blocks > 0,
118+ "inst-1 should have NPU match for tokens_a"
119+ );
120+ assert_eq!(dp0.cpu_blocks, 0, "inst-1 should have no CPU match");
121+ 
122+ // Query with tokens_b — should match inst-2 (CPU) only
123+ let resp = indexer.query("model-b", "t1", &tokens_b, 4).unwrap();
124+ let tenant = &resp.tenants["t1"];
125+ 
126+ let imd2 = &tenant["inst-2"];
127+ let dp2 = &imd2.dp["0"];
128+ assert_eq!(dp2.npu_blocks, 0, "inst-2 should have no NPU match");
129+ assert!(
130+ dp2.cpu_blocks > 0,
131+ "inst-2 should have CPU match for tokens_b"
132+ );
133+}
134+ 
135+#[test]
136+fn test_cpu_continuation_from_hbm_breakpoint() {
137+ let indexer = Indexer::new();
138+ let entry = indexer.get_or_create("model-c", "t1");
139+ 
140+ let tokens: Vec<i64> = (0..12).collect();
141+ let hashes = compute_block_hash_for_seq(&tokens, 4);
142+ assert_eq!(hashes.len(), 3);
143+ 
144+ // HBM holds first two blocks.
145+ let wk_npu = WorkerKey {
146+ instance_id: "inst-1".into(),
147+ backend_id: "inst-1".into(),
148+ dp_rank: 0,
149+ medium: StorageMedium::Npu,
150+ };
151+ entry
152+ .apply_event(
153+ &wk_npu,
154+ &KvCacheEventData::Stored(KvCacheStoreData {
155+ parent_hash: None,
156+ start_position: None,
157+ blocks: vec![
158+ KvCacheStoredBlockData {
159+ block_hash: 100,
160+ tokens_hash: hashes[0].0,
161+ },
162+ KvCacheStoredBlockData {
163+ block_hash: 200,
164+ tokens_hash: hashes[1].0,
165+ },
166+ ],
167+ }),
168+ )
169+ .unwrap();
170+ 
171+ // CPU holds the tail continuing from HBM's last seq_hash=200.
172+ let wk_cpu = WorkerKey {
173+ instance_id: "inst-1".into(),
174+ backend_id: "pool-1".into(),
175+ dp_rank: 0,
176+ medium: StorageMedium::Cpu,
177+ };
178+ entry
179+ .apply_event(
180+ &wk_cpu,
181+ &KvCacheEventData::Stored(KvCacheStoreData {
182+ parent_hash: Some(200),
183+ start_position: None,
184+ blocks: vec![KvCacheStoredBlockData {
185+ block_hash: 300,
186+ tokens_hash: hashes[2].0,
187+ }],
188+ }),
189+ )
190+ .unwrap();
191+ 
192+ let overlap = entry.find_matches(&tokens, 4);
193+ assert_eq!(
194+ overlap.blocks.get(&wk_npu).copied().unwrap_or(0),
195+ 2,
196+ "HBM should match first 2 blocks"
197+ );
198+ assert_eq!(
199+ overlap.blocks.get(&wk_cpu).copied().unwrap_or(0),
200+ 1,
201+ "CPU should continue 1 block from HBM breakpoint"
202+ );
203+ 
204+ // matched_tokens = coverage_end × block_size (HBM 2 + CPU tail 1 → end 3)
205+ let resp = indexer.query("model-c", "t1", &tokens, 4).unwrap();
206+ let dp0 = &resp.tenants["t1"]["inst-1"].dp["0"];
207+ assert_eq!(dp0.npu_blocks, 2);
208+ assert_eq!(dp0.cpu_blocks, 1);
209+ assert_eq!(dp0.disk_blocks, 0);
210+ assert_eq!(dp0.matched_tokens, 3 * 4);
211+ assert_eq!(resp.tenants["t1"]["inst-1"].longest_matched, 12);
212+}
213+ 
214+#[test]
215+fn test_cpu_replica_reported_when_hbm_hits_same_dp() {
216+ let indexer = Indexer::new();
217+ let entry = indexer.get_or_create("model-d", "t1");
218+ 
219+ let tokens: Vec<i64> = vec![1, 2, 3, 4];
220+ let hashes = compute_block_hash_for_seq(&tokens, 4);
221+ assert_eq!(hashes.len(), 1);
222+ 
223+ let wk_npu = WorkerKey {
224+ instance_id: "inst-1".into(),
225+ backend_id: "inst-1".into(),
226+ dp_rank: 0,
227+ medium: StorageMedium::Npu,
228+ };
229+ entry
230+ .apply_event(
231+ &wk_npu,
232+ &KvCacheEventData::Stored(KvCacheStoreData {
233+ parent_hash: None,
234+ start_position: None,
235+ blocks: vec![KvCacheStoredBlockData {
236+ block_hash: 100,
237+ tokens_hash: hashes[0].0,
238+ }],
239+ }),
240+ )
241+ .unwrap();
242+ 
243+ // Same DP also has a full CPU root chain for the same prefix.
244+ let wk_cpu = WorkerKey {
245+ instance_id: "inst-1".into(),
246+ backend_id: "pool-1".into(),
247+ dp_rank: 0,
248+ medium: StorageMedium::Cpu,
249+ };
250+ entry
251+ .apply_event(
252+ &wk_cpu,
253+ &KvCacheEventData::Stored(KvCacheStoreData {
254+ parent_hash: None,
255+ start_position: None,
256+ blocks: vec![KvCacheStoredBlockData {
257+ block_hash: 200,
258+ tokens_hash: hashes[0].0,
259+ }],
260+ }),
261+ )
262+ .unwrap();
263+ 
264+ let resp = indexer.query("model-d", "t1", &tokens, 4).unwrap();
265+ let dp0 = &resp.tenants["t1"]["inst-1"].dp["0"];
266+ assert_eq!(dp0.npu_blocks, 1);
267+ // CPU replica of the same prefix is reported as a real segment; the
268+ // coverage max keeps matched_tokens at the true end (1 block).
269+ assert_eq!(dp0.cpu_blocks, 1);
270+ assert_eq!(dp0.matched_tokens, 4);
271+}
272+ 
273+#[test]
274+fn test_cpu_root_used_when_no_hbm_on_dp() {
275+ let indexer = Indexer::new();
276+ let entry = indexer.get_or_create("model-e", "t1");
277+ 
278+ let tokens: Vec<i64> = vec![9, 8, 7, 6];
279+ let hashes = compute_block_hash_for_seq(&tokens, 4);
280+ 
281+ let wk_cpu = WorkerKey {
282+ instance_id: "inst-cpu-only".into(),
283+ backend_id: "pool-1".into(),
284+ dp_rank: 0,
285+ medium: StorageMedium::Cpu,
286+ };
287+ entry
288+ .apply_event(
289+ &wk_cpu,
290+ &KvCacheEventData::Stored(KvCacheStoreData {
291+ parent_hash: None,
292+ start_position: None,
293+ blocks: vec![KvCacheStoredBlockData {
294+ block_hash: 300,
295+ tokens_hash: hashes[0].0,
296+ }],
297+ }),
298+ )
299+ .unwrap();
300+ 
301+ let resp = indexer.query("model-e", "t1", &tokens, 4).unwrap();
302+ let dp0 = &resp.tenants["t1"]["inst-cpu-only"].dp["0"];
303+ assert_eq!(dp0.npu_blocks, 0);
304+ assert_eq!(dp0.cpu_blocks, 1);
305+ assert_eq!(dp0.matched_tokens, 4);
306+}
307+ 
308+#[test]
309+fn test_disk_continuation_from_cpu_breakpoint() {
310+ let indexer = Indexer::new();
311+ let entry = indexer.get_or_create("model-f", "t1");
312+ 
313+ let tokens: Vec<i64> = (0..12).collect();
314+ let hashes = compute_block_hash_for_seq(&tokens, 4);
315+ assert_eq!(hashes.len(), 3);
316+ 
317+ // HBM: first block
318+ let wk_npu = WorkerKey {
319+ instance_id: "inst-1".into(),
320+ backend_id: "inst-1".into(),
321+ dp_rank: 0,
322+ medium: StorageMedium::Npu,
323+ };
324+ entry
325+ .apply_event(
326+ &wk_npu,
327+ &KvCacheEventData::Stored(KvCacheStoreData {
328+ parent_hash: None,
329+ start_position: None,
330+ blocks: vec![KvCacheStoredBlockData {
331+ block_hash: 100,
332+ tokens_hash: hashes[0].0,
333+ }],
334+ }),
335+ )
336+ .unwrap();
337+ 
338+ // CPU: second block, continuing from HBM seq=100
339+ let wk_cpu = WorkerKey {
340+ instance_id: "inst-1".into(),
341+ backend_id: "pool-cpu".into(),
342+ dp_rank: 0,
343+ medium: StorageMedium::Cpu,
344+ };
345+ entry
346+ .apply_event(
347+ &wk_cpu,
348+ &KvCacheEventData::Stored(KvCacheStoreData {
349+ parent_hash: Some(100),
350+ start_position: None,
351+ blocks: vec![KvCacheStoredBlockData {
352+ block_hash: 200,
353+ tokens_hash: hashes[1].0,
354+ }],
355+ }),
356+ )
357+ .unwrap();
358+ 
359+ // Disk: third block, continuing from CPU seq=200
360+ let wk_disk = WorkerKey {
361+ instance_id: "inst-1".into(),
362+ backend_id: "pool-disk".into(),
363+ dp_rank: 0,
364+ medium: StorageMedium::Disk,
365+ };
366+ entry
367+ .apply_event(
368+ &wk_disk,
369+ &KvCacheEventData::Stored(KvCacheStoreData {
370+ parent_hash: Some(200),
371+ start_position: None,
372+ blocks: vec![KvCacheStoredBlockData {
373+ block_hash: 300,
374+ tokens_hash: hashes[2].0,
375+ }],
376+ }),
377+ )
378+ .unwrap();
379+ 
380+ let resp = indexer.query("model-f", "t1", &tokens, 4).unwrap();
381+ let dp0 = &resp.tenants["t1"]["inst-1"].dp["0"];
382+ assert_eq!(dp0.npu_blocks, 1);
383+ assert_eq!(dp0.cpu_blocks, 1);
384+ assert_eq!(dp0.disk_blocks, 1);
385+ // Coverage end = 3 (consecutive segments), not a double-counted sum.
386+ assert_eq!(dp0.matched_tokens, 3 * 4);
387+}
388+ 
389+#[test]
390+fn test_overlapping_npu_cpu_disk_replicas_do_not_inflate_matched_tokens() {
391+ // Log pattern: same prefix present on NPU + CPU + Disk. Replicas are
392+ // reported as real segment lengths on every tier, but matched_tokens
393+ // must stay within the input prefix (coverage end), never NPU+CPU+Disk.
394+ let indexer = Indexer::new();
395+ let entry = indexer.get_or_create("model-overlap", "t1");
396+ 
397+ let tokens: Vec<i64> = (0..8).collect();
398+ let hashes = compute_block_hash_for_seq(&tokens, 4);
399+ assert_eq!(hashes.len(), 2);
400+ 
401+ let wk_npu = WorkerKey {
402+ instance_id: "inst-1".into(),
403+ backend_id: "inst-1".into(),
404+ dp_rank: 0,
405+ medium: StorageMedium::Npu,
406+ };
407+ entry
408+ .apply_event(
409+ &wk_npu,
410+ &KvCacheEventData::Stored(KvCacheStoreData {
411+ parent_hash: None,
412+ start_position: None,
413+ blocks: vec![
414+ KvCacheStoredBlockData {
415+ block_hash: 100,
416+ tokens_hash: hashes[0].0,
417+ },
418+ KvCacheStoredBlockData {
419+ block_hash: 200,
420+ tokens_hash: hashes[1].0,
421+ },
422+ ],
423+ }),
424+ )
425+ .unwrap();
426+ 
427+ let wk_cpu = WorkerKey {
428+ instance_id: "inst-1".into(),
429+ backend_id: "pool-cpu".into(),
430+ dp_rank: 0,
431+ medium: StorageMedium::Cpu,
432+ };
433+ entry
434+ .apply_event(
435+ &wk_cpu,
436+ &KvCacheEventData::Stored(KvCacheStoreData {
437+ parent_hash: None,
438+ start_position: None,
439+ blocks: vec![
440+ KvCacheStoredBlockData {
441+ block_hash: 100,
442+ tokens_hash: hashes[0].0,
443+ },
444+ KvCacheStoredBlockData {
445+ block_hash: 200,
446+ tokens_hash: hashes[1].0,
447+ },
448+ ],
449+ }),
450+ )
451+ .unwrap();
452+ 
453+ let wk_disk = WorkerKey {
454+ instance_id: "inst-1".into(),
455+ backend_id: "pool-disk".into(),
456+ dp_rank: 0,
457+ medium: StorageMedium::Disk,
458+ };
459+ entry
460+ .apply_event(
461+ &wk_disk,
462+ &KvCacheEventData::Stored(KvCacheStoreData {
463+ parent_hash: None,
464+ start_position: None,
465+ blocks: vec![
466+ KvCacheStoredBlockData {
467+ block_hash: 100,
468+ tokens_hash: hashes[0].0,
469+ },
470+ KvCacheStoredBlockData {
471+ block_hash: 200,
472+ tokens_hash: hashes[1].0,
473+ },
474+ ],
475+ }),
476+ )
477+ .unwrap();
478+ 
479+ let resp = indexer.query("model-overlap", "t1", &tokens, 4).unwrap();
480+ let dp0 = &resp.tenants["t1"]["inst-1"].dp["0"];
481+ assert_eq!(dp0.npu_blocks, 2);
482+ // Replicas on CPU/Disk are reported as real segment lengths (no root
483+ // skip), but matched_tokens stays at the coverage end — never summed.
484+ assert_eq!(dp0.cpu_blocks, 2);
485+ assert_eq!(dp0.disk_blocks, 2);
486+ assert_eq!(dp0.matched_tokens, 2 * 4);
487+ assert!(
488+ dp0.matched_tokens <= tokens.len() as u32,
489+ "matched_tokens {} exceeds input {}",
490+ dp0.matched_tokens,
491+ tokens.len()
492+ );
493+}
494+ 
495+#[test]
496+fn test_shorter_hbm_breakpoint_does_not_overcount_cpu_overlap() {
497+ // Two NPU workers on the same DP with different depths; CPU holds the
498+ // tail after the shorter breakpoint. Coverage must stay at the true
499+ // prefix end (2), not npu_max + cpu_segment.
500+ let indexer = Indexer::new();
501+ let entry = indexer.get_or_create("model-short-break", "t1");
502+ 
503+ let tokens: Vec<i64> = (0..8).collect();
504+ let hashes = compute_block_hash_for_seq(&tokens, 4);
505+ assert_eq!(hashes.len(), 2);
506+ 
507+ let wk_npu_short = WorkerKey {
508+ instance_id: "inst-1".into(),
509+ backend_id: "npu-short".into(),
510+ dp_rank: 0,
511+ medium: StorageMedium::Npu,
512+ };
513+ entry
514+ .apply_event(
515+ &wk_npu_short,
516+ &KvCacheEventData::Stored(KvCacheStoreData {
517+ parent_hash: None,
518+ start_position: None,
519+ blocks: vec![KvCacheStoredBlockData {
520+ block_hash: 100,
521+ tokens_hash: hashes[0].0,
522+ }],
523+ }),
524+ )
525+ .unwrap();
526+ 
527+ let wk_npu_long = WorkerKey {
528+ instance_id: "inst-1".into(),
529+ backend_id: "npu-long".into(),
530+ dp_rank: 0,
531+ medium: StorageMedium::Npu,
532+ };
533+ entry
534+ .apply_event(
535+ &wk_npu_long,
536+ &KvCacheEventData::Stored(KvCacheStoreData {
537+ parent_hash: None,
538+ start_position: None,
539+ blocks: vec![
540+ KvCacheStoredBlockData {
541+ block_hash: 100,
542+ tokens_hash: hashes[0].0,
543+ },
544+ KvCacheStoredBlockData {
545+ block_hash: 200,
546+ tokens_hash: hashes[1].0,
547+ },
548+ ],
549+ }),
550+ )
551+ .unwrap();
552+ 
553+ let wk_cpu = WorkerKey {
554+ instance_id: "inst-1".into(),
555+ backend_id: "pool-cpu".into(),
556+ dp_rank: 0,
557+ medium: StorageMedium::Cpu,
558+ };
559+ entry
560+ .apply_event(
561+ &wk_cpu,
562+ &KvCacheEventData::Stored(KvCacheStoreData {
563+ parent_hash: Some(100),
564+ start_position: None,
565+ blocks: vec![KvCacheStoredBlockData {
566+ block_hash: 200,
567+ tokens_hash: hashes[1].0,
568+ }],
569+ }),
570+ )
571+ .unwrap();
572+ 
573+ let resp = indexer
574+ .query("model-short-break", "t1", &tokens, 4)
575+ .unwrap();
576+ let dp0 = &resp.tenants["t1"]["inst-1"].dp["0"];
577+ assert_eq!(dp0.npu_blocks, 2);
578+ // The CPU worker genuinely owns the chained block (100→200 = position 1),
579+ // so it scores a 1-block continuation segment. The farthest HBM
580+ // breakpoint (end_pos=2) sits at the sequence end and cannot continue —
581+ // but coverage end is still 2, so matched_tokens stays within input.
582+ assert_eq!(dp0.cpu_blocks, 1);
583+ assert_eq!(dp0.matched_tokens, 2 * 4);
584+}
585+ 
586+#[test]
587+fn test_disk_continuation_from_hbm_when_cpu_miss() {
588+ // vLLM lookup: after NPU hit, Disk can hit even if CPU miss (then promote).
589+ let indexer = Indexer::new();
590+ let entry = indexer.get_or_create("model-h", "t1");
591+ 
592+ let tokens: Vec<i64> = (0..8).collect();
593+ let hashes = compute_block_hash_for_seq(&tokens, 4);
594+ assert_eq!(hashes.len(), 2);
595+ 
596+ let wk_npu = WorkerKey {
597+ instance_id: "inst-1".into(),
598+ backend_id: "inst-1".into(),
599+ dp_rank: 0,
600+ medium: StorageMedium::Npu,
601+ };
602+ entry
603+ .apply_event(
604+ &wk_npu,
605+ &KvCacheEventData::Stored(KvCacheStoreData {
606+ parent_hash: None,
607+ start_position: None,
608+ blocks: vec![KvCacheStoredBlockData {
609+ block_hash: 100,
610+ tokens_hash: hashes[0].0,
611+ }],
612+ }),
613+ )
614+ .unwrap();
615+ 
616+ // Disk tail continues from HBM; CPU has nothing.
617+ let wk_disk = WorkerKey {
618+ instance_id: "inst-1".into(),
619+ backend_id: "pool-disk".into(),
620+ dp_rank: 0,
621+ medium: StorageMedium::Disk,
622+ };
623+ entry
624+ .apply_event(
625+ &wk_disk,
626+ &KvCacheEventData::Stored(KvCacheStoreData {
627+ parent_hash: Some(100),
628+ start_position: None,
629+ blocks: vec![KvCacheStoredBlockData {
630+ block_hash: 300,
631+ tokens_hash: hashes[1].0,
632+ }],
633+ }),
634+ )
635+ .unwrap();
636+ 
637+ let resp = indexer.query("model-h", "t1", &tokens, 4).unwrap();
638+ let dp0 = &resp.tenants["t1"]["inst-1"].dp["0"];
639+ assert_eq!(dp0.npu_blocks, 1);
640+ assert_eq!(dp0.cpu_blocks, 0);
641+ assert_eq!(dp0.disk_blocks, 1);
642+ assert_eq!(dp0.matched_tokens, 2 * 4); // coverage end after NPU + disk tail
643+}
644+ 
645+#[test]
646+fn test_disk_replica_reported_when_cpu_hits_same_dp() {
647+ let indexer = Indexer::new();
648+ let entry = indexer.get_or_create("model-g", "t1");
649+ 
650+ let tokens: Vec<i64> = vec![1, 2, 3, 4];
651+ let hashes = compute_block_hash_for_seq(&tokens, 4);
652+ 
653+ let wk_cpu = WorkerKey {
654+ instance_id: "inst-1".into(),
655+ backend_id: "pool-cpu".into(),
656+ dp_rank: 0,
657+ medium: StorageMedium::Cpu,
658+ };
659+ entry
660+ .apply_event(
661+ &wk_cpu,
662+ &KvCacheEventData::Stored(KvCacheStoreData {
663+ parent_hash: None,
664+ start_position: None,
665+ blocks: vec![KvCacheStoredBlockData {
666+ block_hash: 100,
667+ tokens_hash: hashes[0].0,
668+ }],
669+ }),
670+ )
671+ .unwrap();
672+ 
673+ // Same DP also has a Disk root chain — reported as a real segment.
674+ let wk_disk = WorkerKey {
675+ instance_id: "inst-1".into(),
676+ backend_id: "pool-disk".into(),
677+ dp_rank: 0,
678+ medium: StorageMedium::Disk,
679+ };
680+ entry
681+ .apply_event(
682+ &wk_disk,
683+ &KvCacheEventData::Stored(KvCacheStoreData {
684+ parent_hash: None,
685+ start_position: None,
686+ blocks: vec![KvCacheStoredBlockData {
687+ block_hash: 200,
688+ tokens_hash: hashes[0].0,
689+ }],
690+ }),
691+ )
692+ .unwrap();
693+ 
694+ let resp = indexer.query("model-g", "t1", &tokens, 4).unwrap();
695+ let dp0 = &resp.tenants["t1"]["inst-1"].dp["0"];
696+ assert_eq!(dp0.cpu_blocks, 1);
697+ // Disk replica is reported as a real segment; coverage max keeps
698+ // matched_tokens at the true end (1 block).
699+ assert_eq!(dp0.disk_blocks, 1);
700+ assert_eq!(dp0.matched_tokens, 4);
701+}
702+ 
703+#[test]
704+fn test_disk_replica_reported_when_hbm_hits_same_dp() {
705+ let indexer = Indexer::new();
706+ let entry = indexer.get_or_create("model-i", "t1");
707+ 
708+ let tokens: Vec<i64> = vec![1, 2, 3, 4];
709+ let hashes = compute_block_hash_for_seq(&tokens, 4);
710+ 
711+ let wk_npu = WorkerKey {
712+ instance_id: "inst-1".into(),
713+ backend_id: "inst-1".into(),
714+ dp_rank: 0,
715+ medium: StorageMedium::Npu,
716+ };
717+ entry
718+ .apply_event(
719+ &wk_npu,
720+ &KvCacheEventData::Stored(KvCacheStoreData {
721+ parent_hash: None,
722+ start_position: None,
723+ blocks: vec![KvCacheStoredBlockData {
724+ block_hash: 100,
725+ tokens_hash: hashes[0].0,
726+ }],
727+ }),
728+ )
729+ .unwrap();
730+ 
731+ let wk_disk = WorkerKey {
732+ instance_id: "inst-1".into(),
733+ backend_id: "pool-disk".into(),
734+ dp_rank: 0,
735+ medium: StorageMedium::Disk,
736+ };
737+ entry
738+ .apply_event(
739+ &wk_disk,
740+ &KvCacheEventData::Stored(KvCacheStoreData {
741+ parent_hash: None,
742+ start_position: None,
743+ blocks: vec![KvCacheStoredBlockData {
744+ block_hash: 200,
745+ tokens_hash: hashes[0].0,
746+ }],
747+ }),
748+ )
749+ .unwrap();
750+ 
751+ let resp = indexer.query("model-i", "t1", &tokens, 4).unwrap();
752+ let dp0 = &resp.tenants["t1"]["inst-1"].dp["0"];
753+ assert_eq!(dp0.npu_blocks, 1);
754+ // Disk replica of the same prefix is reported as a real segment;
755+ // coverage max keeps matched_tokens at the true end (1 block).
756+ assert_eq!(dp0.disk_blocks, 1);
757+ assert_eq!(dp0.matched_tokens, 4);
758+}
759+ 
760+/// The documented "longer lower-tier replica" case: NPU holds only the
761+/// first block while Disk holds the full prefix. The old root-skip made
762+/// this silently under-report (disk_blocks=0, matched=1 block); the
763+/// unconditional root walk must surface the Disk copy and extend the
764+/// coverage end to the true prefix length.
765+#[test]
766+fn test_lower_tier_longer_replica_extends_coverage() {
767+ let indexer = Indexer::new();
768+ let entry = indexer.get_or_create("model-long-replica", "t1");
769+ 
770+ let tokens: Vec<i64> = (0..12).collect();
771+ let hashes = compute_block_hash_for_seq(&tokens, 4);
772+ assert_eq!(hashes.len(), 3);
773+ 
774+ // NPU: only the first block.
775+ let wk_npu = WorkerKey {
776+ instance_id: "inst-1".into(),
777+ backend_id: "inst-1".into(),
778+ dp_rank: 0,
779+ medium: StorageMedium::Npu,
780+ };
781+ entry
782+ .apply_event(
783+ &wk_npu,
784+ &KvCacheEventData::Stored(KvCacheStoreData {
785+ parent_hash: None,
786+ start_position: None,
787+ blocks: vec![KvCacheStoredBlockData {
788+ block_hash: 100,
789+ tokens_hash: hashes[0].0,
790+ }],
791+ }),
792+ )
793+ .unwrap();
794+ 
795+ // Disk: full 3-block root chain.
796+ let wk_disk = WorkerKey {
797+ instance_id: "inst-1".into(),
798+ backend_id: "pool-disk".into(),
799+ dp_rank: 0,
800+ medium: StorageMedium::Disk,
801+ };
802+ entry
803+ .apply_event(
804+ &wk_disk,
805+ &KvCacheEventData::Stored(KvCacheStoreData {
806+ parent_hash: None,
807+ start_position: None,
808+ blocks: vec![
809+ KvCacheStoredBlockData {
810+ block_hash: 200,
811+ tokens_hash: hashes[0].0,
812+ },
813+ KvCacheStoredBlockData {
814+ block_hash: 201,
815+ tokens_hash: hashes[1].0,
816+ },
817+ KvCacheStoredBlockData {
818+ block_hash: 202,
819+ tokens_hash: hashes[2].0,
820+ },
821+ ],
822+ }),
823+ )
824+ .unwrap();
825+ 
826+ let resp = indexer
827+ .query("model-long-replica", "t1", &tokens, 4)
828+ .unwrap();
829+ let dp0 = &resp.tenants["t1"]["inst-1"].dp["0"];
830+ assert_eq!(dp0.npu_blocks, 1);
831+ assert_eq!(
832+ dp0.disk_blocks, 3,
833+ "Disk replica of the full prefix must be reported"
834+ );
835+ assert_eq!(
836+ dp0.matched_tokens,
837+ 3 * 4,
838+ "coverage must extend to the longer replica, not the NPU's 1 block"
839+ );
840+ assert!(dp0.matched_tokens <= tokens.len() as u32);
841+}
842+ 
843+#[test]
844+fn test_no_indexer_error() {
845+ let indexer = Indexer::new();
846+ let err = indexer.query("no-such-model", "default", &[1, 2, 3, 4], 4);
847+ assert!(err.is_err());
848+ assert!(matches!(
849+ err.unwrap_err(),
850+ KvConductorError::NoIndexer { .. }
851+ ));
852+}
Mmotor/kv_conductor/src/lib.rs+2-1
@@ -21,6 +21,7 @@ pub mod error;
21pub mod events;21pub mod events;
22pub mod hashing;22pub mod hashing;
23pub mod indexer;23pub mod indexer;
24+pub mod lower_tier;
24pub mod protocols;25pub mod protocols;
25pub mod registry;26pub mod registry;
26pub mod server;27pub mod server;
@@ -36,7 +37,7 @@ pub use indexer::Indexer;
36pub use protocols::{37pub use protocols::{
37 DpRank, HbmIpIndex, InstanceId, InstanceMatchData, KvCacheEvent, KvCacheEventData,38 DpRank, HbmIpIndex, InstanceId, InstanceMatchData, KvCacheEvent, KvCacheEventData,
38 KvCacheStoreData, KvCacheStoredBlockData, KvEventBatch, KvEventWirePayload, LocalBlockHash,39 KvCacheStoreData, KvCacheStoredBlockData, KvEventBatch, KvEventWirePayload, LocalBlockHash,
39- OverlapScores, QueryByHashRequest, QueryRequest, QueryResponse, RegisterRequest,40+ OverlapBlocks, QueryByHashRequest, QueryRequest, QueryResponse, RegisterRequest,
40 SequenceBlockHash, StorageMedium, UnregisterRequest, WorkerKey,41 SequenceBlockHash, StorageMedium, UnregisterRequest, WorkerKey,
41};42};
42pub use registry::WorkerRegistry;43pub use registry::WorkerRegistry;
Amotor/kv_conductor/src/lower_tier.rs+515-0
@@ -0,0 +1,515 @@
1+// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+// SPDX-FileCopyrightText: Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.
3+// SPDX-License-Identifier: Apache-2.0
4+//
5+// This file is a Derivative Work of NVIDIA Dynamo kv-router
6+// (https://github.com/ai-dynamo/dynamo), originally licensed under the
7+// Apache License, Version 2.0. Upstream source path:
8+// lib/kv-router/src/indexer/lower_tier.rs
9+//
10+// You may obtain a copy of the Apache License at:
11+// http://www.apache.org/licenses/LICENSE-2.0
12+// Local copy: licenses/Apache-2.0.txt
13+// Attribution: THIRD_PARTY_NOTICES.md
14+//
15+// Modified by Huawei Technologies Co., Ltd. for MindIE-PyMotor KV Conductor
16+// (RwLock + per-worker reverse index, ContiguousHit API, WorkerKey/medium
17+// integration). Huawei modifications are also available under Mulan PSL v2
18+// (http://license.coscl.org.cn/MulanPSL2). Redistribution of this file must
19+// still comply with Apache License 2.0.
20+ 
21+//! Lower-tier (CPU / Disk) continuation-edge index.
22+//!
23+//! Derived from NVIDIA Dynamo kv-router `LowerTierIndexer`
24+//! (`lib/kv-router/src/indexer/lower_tier.rs`, Apache-2.0). See
25+//! `THIRD_PARTY_NOTICES.md`.
26+//!
27+//! Stores worker ownership over shared continuation edges:
28+//! ``(parent_sequence_hash, local_hash) -> child_sequence_hash``.
29+//!
30+//! Unlike the HBM radix tree, this index does **not** score from root by
31+//! default. Queries continue from caller-provided per-worker continuation
32+//! points (HBM → CPU; max(HBM, CPU) → Disk) and count how many
33+//! **consecutive** lower-tier blocks are present. The caller also issues an
34+//! unconditional root-walk candidate per worker owning the first edge, so a
35+//! longer full replica on this tier is never hidden by a shorter upstream
36+//! hit; the walk keeps the candidate with the farthest absolute end.
37+ 
38+use parking_lot::RwLock;
39+use rustc_hash::{FxHashMap, FxHashSet};
40+ 
41+use crate::protocols::{KvCacheStoreData, LocalBlockHash, SequenceBlockHash, WorkerKey};
42+ 
43+type WorkerSet = FxHashSet<WorkerKey>;
44+ 
45+/// Edge key in the continuation graph.
46+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
47+struct TransitionKey {
48+ parent_hash: Option<SequenceBlockHash>,
49+ local_hash: LocalBlockHash,
50+}
51+ 
52+#[derive(Debug, Clone)]
53+enum EdgeOwnersEntry {
54+ Single {
55+ child_hash: SequenceBlockHash,
56+ owner: WorkerKey,
57+ },
58+ Multi {
59+ child_hash: SequenceBlockHash,
60+ owners: WorkerSet,
61+ },
62+}
63+ 
64+impl EdgeOwnersEntry {
65+ fn new(child_hash: SequenceBlockHash, owner: WorkerKey) -> Self {
66+ Self::Single { child_hash, owner }
67+ }
68+ 
69+ fn child_hash(&self) -> SequenceBlockHash {
70+ match self {
71+ Self::Single { child_hash, .. } | Self::Multi { child_hash, .. } => *child_hash,
72+ }
73+ }
74+ 
75+ /// Insert `owner` for this edge. Returns `false` if `child_hash` conflicts
76+ /// with the existing mapping (first-writer wins).
77+ fn insert(&mut self, child_hash: SequenceBlockHash, owner: WorkerKey) -> bool {
78+ match self {
79+ Self::Single {
80+ child_hash: existing_hash,
81+ owner: existing_owner,
82+ } => {
83+ if *existing_hash != child_hash {
84+ return false;
85+ }
86+ if *existing_owner == owner {
87+ return true;
88+ }
89+ let mut owners = WorkerSet::default();
90+ owners.insert(existing_owner.clone());
91+ owners.insert(owner);
92+ *self = Self::Multi { child_hash, owners };
93+ true
94+ }
95+ Self::Multi {
96+ child_hash: existing_hash,
97+ owners,
98+ } => {
99+ if *existing_hash != child_hash {
100+ return false;
101+ }
102+ owners.insert(owner);
103+ true
104+ }
105+ }
106+ }
107+ 
108+ /// Remove `owner`. Returns `true` if the edge should be deleted.
109+ fn remove(&mut self, owner: &WorkerKey) -> bool {
110+ match self {
111+ Self::Single {
112+ owner: existing_owner,
113+ ..
114+ } => existing_owner == owner,
115+ Self::Multi { child_hash, owners } => {
116+ if !owners.remove(owner) {
117+ return false;
118+ }
119+ if owners.is_empty() {
120+ return true;
121+ }
122+ if owners.len() == 1 {
123+ let remaining = owners.iter().next().cloned().unwrap();
124+ *self = Self::Single {
125+ child_hash: *child_hash,
126+ owner: remaining,
127+ };
128+ }
129+ false
130+ }
131+ }
132+ }
133+ 
134+ fn contains(&self, owner: &WorkerKey) -> bool {
135+ match self {
136+ Self::Single {
137+ owner: existing_owner,
138+ ..
139+ } => existing_owner == owner,
140+ Self::Multi { owners, .. } => owners.contains(owner),
141+ }
142+ }
143+ 
144+ fn collect_workers(&self) -> Vec<WorkerKey> {
145+ match self {
146+ Self::Single { owner, .. } => vec![owner.clone()],
147+ Self::Multi { owners, .. } => owners.iter().cloned().collect(),
148+ }
149+ }
150+}
151+ 
152+/// Where a lower-tier walk should resume for one worker.
153+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
154+pub struct LowerTierContinuation {
155+ pub start_pos: usize,
156+ pub last_matched_hash: Option<SequenceBlockHash>,
157+}
158+ 
159+impl LowerTierContinuation {
160+ pub fn new(start_pos: usize, last_matched_hash: SequenceBlockHash) -> Self {
161+ Self {
162+ start_pos,
163+ last_matched_hash: Some(last_matched_hash),
164+ }
165+ }
166+ 
167+ pub fn from_root(start_pos: usize) -> Self {
168+ Self {
169+ start_pos,
170+ last_matched_hash: None,
171+ }
172+ }
173+}
174+ 
175+/// Result of a contiguous lower-tier walk for one worker.
176+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
177+pub struct ContiguousHit {
178+ /// Number of blocks matched from ``start_pos``.
179+ pub count: usize,
180+ /// Absolute start index in the query hash sequence.
181+ pub start_pos: usize,
182+ /// Sequence hash of the last matched block (for the next tier).
183+ pub last_matched_hash: Option<SequenceBlockHash>,
184+}
185+ 
186+impl ContiguousHit {
187+ /// Absolute end index (exclusive) — next tier continues here.
188+ pub fn end_pos(&self) -> usize {
189+ self.start_pos.saturating_add(self.count)
190+ }
191+}
192+ 
193+/// Continuation-edge index for one lower-tier medium (CPU or Disk).
194+#[derive(Debug, Default)]
195+pub struct LowerTierIndexer {
196+ edges: RwLock<FxHashMap<TransitionKey, EdgeOwnersEntry>>,
197+ /// Per-worker reverse lookup: ``block_hash ->TransitionKey`` for O(1) remove.
198+ worker_blocks: RwLock<FxHashMap<WorkerKey, FxHashMap<SequenceBlockHash, TransitionKey>>>,
199+}
200+ 
201+impl LowerTierIndexer {
202+ pub fn new() -> Self {
203+ Self::default()
204+ }
205+ 
206+ /// Workers owning the root edge for ``local_hash`` (parent = None).
207+ pub fn root_workers(&self, local_hash: LocalBlockHash) -> Vec<WorkerKey> {
208+ self.edge_owners(None, local_hash)
209+ }
210+ 
211+ /// Workers owning edge ``(parent_hash, local_hash)``.
212+ pub fn edge_owners(
213+ &self,
214+ parent_hash: Option<SequenceBlockHash>,
215+ local_hash: LocalBlockHash,
216+ ) -> Vec<WorkerKey> {
217+ let key = TransitionKey {
218+ parent_hash,
219+ local_hash,
220+ };
221+ self.edges
222+ .read()
223+ .get(&key)
224+ .map(|e| e.collect_workers())
225+ .unwrap_or_default()
226+ }
227+ 
228+ /// Insert a stored chain as continuation edges.
229+ pub fn store_blocks(&self, worker: &WorkerKey, store_data: &KvCacheStoreData) {
230+ let mut parent_hash = store_data.parent_hash.map(SequenceBlockHash);
231+ let mut worker_blocks = self.worker_blocks.write();
232+ let worker_map = worker_blocks.entry(worker.clone()).or_default();
233+ let mut edges = self.edges.write();
234+ 
235+ for block in &store_data.blocks {
236+ let child = SequenceBlockHash(block.block_hash);
237+ let key = TransitionKey {
238+ parent_hash,
239+ local_hash: LocalBlockHash(block.tokens_hash),
240+ };
241+ 
242+ // Conflicting reverse mapping for the same block_hash ->stop chain.
243+ if worker_map
244+ .get(&child)
245+ .is_some_and(|existing| *existing != key)
246+ {
247+ break;
248+ }
249+ 
250+ let inserted = match edges.get_mut(&key) {
251+ Some(edge) => edge.insert(child, worker.clone()),
252+ None => {
253+ edges.insert(key, EdgeOwnersEntry::new(child, worker.clone()));
254+ true
255+ }
256+ };
257+ 
258+ if !inserted {
259+ break;
260+ }
261+ 
262+ worker_map.insert(child, key);
263+ parent_hash = Some(child);
264+ }
265+ }
266+ 
267+ /// Remove blocks by engine sequence hash.
268+ pub fn remove_blocks(&self, worker: &WorkerKey, block_hashes: &[u64]) {
269+ let mut worker_blocks = self.worker_blocks.write();
270+ let Some(worker_map) = worker_blocks.get_mut(worker) else {
271+ return;
272+ };
273+ let mut edges = self.edges.write();
274+ 
275+ for &h in block_hashes {
276+ let seq = SequenceBlockHash(h);
277+ let Some(key) = worker_map.remove(&seq) else {
278+ continue;
279+ };
280+ if let Some(edge) = edges.get_mut(&key) {
281+ if edge.remove(worker) {
282+ edges.remove(&key);
283+ }
284+ }
285+ }
286+ 
287+ if worker_map.is_empty() {
288+ worker_blocks.remove(worker);
289+ }
290+ }
291+ 
292+ /// Drop all edges owned by ``worker``.
293+ pub fn clear_worker(&self, worker: &WorkerKey) {
294+ let mut worker_blocks = self.worker_blocks.write();
295+ let Some(worker_map) = worker_blocks.remove(worker) else {
296+ return;
297+ };
298+ let mut edges = self.edges.write();
299+ for (_, key) in worker_map {
300+ if let Some(edge) = edges.get_mut(&key) {
301+ if edge.remove(worker) {
302+ edges.remove(&key);
303+ }
304+ }
305+ }
306+ }
307+ 
308+ /// Look up `(parent_hash, tokens_hash)` for an engine `block_hash` owned
309+ /// by any worker in this tier.
310+ ///
311+ /// Used when a later pool medium (e.g. Disk) confirms a block that was
312+ /// already indexed on another lower tier (e.g. CPU): reuse the content
313+ /// mapping without requiring a fresh engine offload event.
314+ pub fn lookup_block(&self, block_hash: u64) -> Option<(Option<u64>, u64)> {
315+ let seq = SequenceBlockHash(block_hash);
316+ let worker_blocks = self.worker_blocks.read();
317+ for worker_map in worker_blocks.values() {
318+ if let Some(key) = worker_map.get(&seq) {
319+ return Some((key.parent_hash.map(|h| h.0), key.local_hash.0));
320+ }
321+ }
322+ None
323+ }
324+ 
325+ /// Whether any worker currently owns ``block_hash``.
326+ pub fn contains_block(&self, block_hash: u64) -> bool {
327+ let seq = SequenceBlockHash(block_hash);
328+ self.worker_blocks
329+ .read()
330+ .values()
331+ .any(|m| m.contains_key(&seq))
332+ }
333+ 
334+ /// Number of blocks tracked for ``worker``.
335+ pub fn worker_block_count(&self, worker: &WorkerKey) -> usize {
336+ self.worker_blocks
337+ .read()
338+ .get(worker)
339+ .map(|m| m.len())
340+ .unwrap_or(0)
341+ }
342+ 
343+ /// Total blocks across all workers (sum of reverse-lookup sizes).
344+ pub fn total_blocks(&self) -> usize {
345+ self.worker_blocks.read().values().map(|m| m.len()).sum()
346+ }
347+ 
348+ /// All workers that currently own at least one edge.
349+ pub fn worker_keys(&self) -> Vec<WorkerKey> {
350+ self.worker_blocks.read().keys().cloned().collect()
351+ }
352+ 
353+ pub fn is_empty(&self) -> bool {
354+ self.worker_blocks.read().is_empty()
355+ }
356+ 
357+ /// For each worker, walk contiguous lower-tier hits from its continuations.
358+ ///
359+ /// A worker may have several candidate continuations (a root walk plus one
360+ /// or more upstream-breakpoint continuations). Each candidate is walked
361+ /// independently and the one with the **farthest absolute end** wins —
362+ /// coverage semantics stay correct no matter which candidate is longer.
363+ pub fn query_contiguous_hits(
364+ &self,
365+ local_hashes: &[LocalBlockHash],
366+ continuations: &FxHashMap<WorkerKey, Vec<LowerTierContinuation>>,
367+ ) -> FxHashMap<WorkerKey, ContiguousHit> {
368+ let mut hits = FxHashMap::default();
369+ let edges = self.edges.read();
370+ 
371+ for (worker, conts) in continuations {
372+ let mut best: Option<ContiguousHit> = None;
373+ for cont in conts {
374+ let mut cur_pos = cont.start_pos;
375+ let mut cur_hash = cont.last_matched_hash;
376+ let start = cur_pos;
377+ 
378+ while cur_pos < local_hashes.len() {
379+ let key = TransitionKey {
380+ parent_hash: cur_hash,
381+ local_hash: local_hashes[cur_pos],
382+ };
383+ let Some(edge) = edges.get(&key) else {
384+ break;
385+ };
386+ if !edge.contains(worker) {
387+ break;
388+ }
389+ cur_hash = Some(edge.child_hash());
390+ cur_pos += 1;
391+ }
392+ 
393+ let hit = ContiguousHit {
394+ count: cur_pos.saturating_sub(start),
395+ start_pos: start,
396+ last_matched_hash: if cur_pos > start { cur_hash } else { None },
397+ };
398+ // Keep the candidate with the farthest absolute end (ties:
399+ // later candidate wins — same end implies the same last hash).
400+ best = match best {
401+ Some(b) if hit.end_pos() >= b.end_pos() => Some(hit),
402+ Some(b) => Some(b),
403+ None => Some(hit),
404+ };
405+ }
406+ if let Some(b) = best {
407+ hits.insert(worker.clone(), b);
408+ }
409+ }
410+ 
411+ hits
412+ }
413+}
414+ 
415+#[cfg(test)]
416+mod tests {
417+ use super::*;
418+ use crate::protocols::{KvCacheStoredBlockData, StorageMedium};
419+ 
420+ fn worker(id: &str) -> WorkerKey {
421+ WorkerKey {
422+ instance_id: id.into(),
423+ backend_id: id.into(),
424+ dp_rank: 0,
425+ medium: StorageMedium::Cpu,
426+ }
427+ }
428+ 
429+ fn store(parent: Option<u64>, blocks: &[(u64, u64)]) -> KvCacheStoreData {
430+ KvCacheStoreData {
431+ parent_hash: parent,
432+ start_position: None,
433+ blocks: blocks
434+ .iter()
435+ .map(|&(bh, th)| KvCacheStoredBlockData {
436+ block_hash: bh,
437+ tokens_hash: th,
438+ })
439+ .collect(),
440+ }
441+ }
442+ 
443+ #[test]
444+ fn root_chain_full_match() {
445+ let idx = LowerTierIndexer::new();
446+ let w = worker("w1");
447+ idx.store_blocks(&w, &store(None, &[(101, 11), (102, 12)]));
448+ 
449+ let mut conts = FxHashMap::default();
450+ conts.insert(w.clone(), vec![LowerTierContinuation::from_root(0)]);
451+ let hits = idx.query_contiguous_hits(&[LocalBlockHash(11), LocalBlockHash(12)], &conts);
452+ assert_eq!(hits.get(&w).map(|h| h.count), Some(2));
453+ }
454+ 
455+ #[test]
456+ fn mid_chain_continuation_from_parent() {
457+ let idx = LowerTierIndexer::new();
458+ let w = worker("w1");
459+ // Tail only: parent=999, then local 21,22
460+ idx.store_blocks(&w, &store(Some(999), &[(201, 21), (202, 22)]));
461+ 
462+ let mut conts = FxHashMap::default();
463+ conts.insert(
464+ w.clone(),
465+ vec![LowerTierContinuation::new(2, SequenceBlockHash(999))],
466+ );
467+ let query = [
468+ LocalBlockHash(1),
469+ LocalBlockHash(2),
470+ LocalBlockHash(21),
471+ LocalBlockHash(22),
472+ ];
473+ let hits = idx.query_contiguous_hits(&query, &conts);
474+ assert_eq!(hits.get(&w).map(|h| h.count), Some(2));
475+ assert_eq!(
476+ hits.get(&w).and_then(|h| h.last_matched_hash),
477+ Some(SequenceBlockHash(202))
478+ );
479+ assert_eq!(hits.get(&w).map(|h| h.end_pos()), Some(4));
480+ }
481+ 
482+ #[test]
483+ fn remove_breaks_contiguous_walk() {
484+ let idx = LowerTierIndexer::new();
485+ let w = worker("w1");
486+ idx.store_blocks(&w, &store(None, &[(101, 11), (102, 12), (103, 13)]));
487+ idx.remove_blocks(&w, &[102]);
488+ 
489+ let mut conts = FxHashMap::default();
490+ conts.insert(w.clone(), vec![LowerTierContinuation::from_root(0)]);
491+ let hits = idx.query_contiguous_hits(
492+ &[LocalBlockHash(11), LocalBlockHash(12), LocalBlockHash(13)],
493+ &conts,
494+ );
495+ // First edge remains; walk stops at missing middle edge.
496+ assert_eq!(hits.get(&w).map(|h| h.count), Some(1));
497+ }
498+ 
499+ #[test]
500+ fn shared_edge_remove_preserves_other_owner() {
501+ let idx = LowerTierIndexer::new();
502+ let a = worker("a");
503+ let b = worker("b");
504+ idx.store_blocks(&a, &store(None, &[(101, 11), (102, 12)]));
505+ idx.store_blocks(&b, &store(None, &[(101, 11), (102, 12)]));
506+ idx.remove_blocks(&a, &[101, 102]);
507+ 
508+ let mut conts = FxHashMap::default();
509+ conts.insert(a.clone(), vec![LowerTierContinuation::from_root(0)]);
510+ conts.insert(b.clone(), vec![LowerTierContinuation::from_root(0)]);
511+ let hits = idx.query_contiguous_hits(&[LocalBlockHash(11), LocalBlockHash(12)], &conts);
512+ assert_eq!(hits.get(&a).map(|h| h.count), Some(0));
513+ assert_eq!(hits.get(&b).map(|h| h.count), Some(2));
514+ }
515+}
jason lyu
jason lyujason lyu19 天前

根起步仅 worker 拥第一条边时能继续，若根边不存在但后续边存在，查询会漏，请确认是否预期。

likedislike
G
ganglv
18 天前 评论：
Mmotor/kv_conductor/src/main.rs+2-27
@@ -21,7 +21,6 @@ use clap::Parser;
21use tracing_subscriber::fmt::time::OffsetTime;21use tracing_subscriber::fmt::time::OffsetTime;
22use tracing_subscriber::EnvFilter;22use tracing_subscriber::EnvFilter;
23 23 
24-use kv_conductor::protocols::ScoringConfig;
25use kv_conductor::registry::WorkerRegistry;24use kv_conductor::registry::WorkerRegistry;
26use kv_conductor::server::{create_router, AppState};25use kv_conductor::server::{create_router, AppState};
27 26 
@@ -37,18 +36,6 @@ struct Cli {
37 /// Port to listen on36 /// Port to listen on
38 #[arg(long, short, default_value = "13333")]37 #[arg(long, short, default_value = "13333")]
39 port: u16,38 port: u16,
40- 
41- /// Score per matched HBM/XPU block
42- #[arg(long, default_value = "3")]
43- hbm_weight: u32,
44- 
45- /// Score per matched CPU block
46- #[arg(long, default_value = "2")]
47- cpu_weight: u32,
48- 
49- /// Score per matched disk block
50- #[arg(long, default_value = "1")]
51- disk_weight: u32,
52}39}
53 40 
54#[tokio::main]41#[tokio::main]
@@ -70,20 +57,8 @@ async fn main() {
70 let host: IpAddr = cli.host.parse().expect("invalid host address");57 let host: IpAddr = cli.host.parse().expect("invalid host address");
71 let addr = SocketAddr::new(host, cli.port);58 let addr = SocketAddr::new(host, cli.port);
72 59 
73- let scoring = ScoringConfig {60+ let registry = Arc::new(WorkerRegistry::new());
74- hbm_weight: cli.hbm_weight,61+ let state = AppState { registry };
75- cpu_weight: cli.cpu_weight,
76- disk_weight: cli.disk_weight,
77- };
78- tracing::info!(
79- hbm_weight = scoring.hbm_weight,
80- cpu_weight = scoring.cpu_weight,
81- disk_weight = scoring.disk_weight,
82- "scoring config"
83- );
84- 
85- let registry = Arc::new(WorkerRegistry::new(scoring.clone()));
86- let state = AppState { registry, scoring };
87 let router = create_router(state);62 let router = create_router(state);
88 63 
89 tracing::info!("KV conductor starting on {}", addr);64 tracing::info!("KV conductor starting on {}", addr);
Mmotor/kv_conductor/src/protocols.rs+232-141
@@ -1,16 +1,28 @@
1-// Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.1+// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2-// MindIE is licensed under Mulan PSL v2.2+// SPDX-FileCopyrightText: Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.
3-// You can use this software according to the terms and conditions of the Mulan PSL v2.3+//
4-// You may obtain a copy of Mulan PSL v2 at:4+// Portions of this file (hash newtypes, KV cache event / store payloads, and
5-// http://license.coscl.org.cn/MulanPSL25+// overlap-match result types) are a Derivative Work of NVIDIA Dynamo
6-// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,6+// kv-router lib/kv-router/src/protocols.rs, licensed under Apache-2.0.
7-// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,7+// Upstream project: https://github.com/ai-dynamo/dynamo
8-// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.8+// SPDX-License-Identifier: Apache-2.0
9-// See the Mulan PSL v2 for more details.9+//
10+// You may obtain a copy of the Apache License at:
11+// http://www.apache.org/licenses/LICENSE-2.0
12+// Local copy: licenses/Apache-2.0.txt
13+// Attribution: THIRD_PARTY_NOTICES.md
14+//
15+// MindIE HTTP API types (register / query / health, WorkerKey, StorageMedium
16+// parsing, HbmIpIndex, etc.) and other Huawei modifications are also
17+// available under Mulan PSL v2 (http://license.coscl.org.cn/MulanPSL2).
18+// Redistribution of the Dynamo-derived portions must still comply with
19+// Apache License 2.0.
10 20 
11//! Protocol types for the KV conductor service.21//! Protocol types for the KV conductor service.
12//!22//!
13-//! These types define the HTTP API contract, compatible with the Python23+//! Hash / KV-event payload types are derived from NVIDIA Dynamo kv-router
24+//! `lib/kv-router/src/protocols.rs` (Apache-2.0); see `THIRD_PARTY_NOTICES.md`.
25+//! HTTP API types are the MindIE conductor contract used by Python
14//! `ConductorApiClient` in `motor/coordinator/api_client/`.26//! `ConductorApiClient` in `motor/coordinator/api_client/`.
15 27 
16use std::collections::HashMap;28use std::collections::HashMap;
@@ -21,6 +33,8 @@ use rustc_hash::FxHashMap;
21use serde::{Deserialize, Serialize};33use serde::{Deserialize, Serialize};
22use tracing;34use tracing;
23 35 
36+use crate::hashing::compute_block_hash_for_seq;
37+ 
24// ---------------------------------------------------------------------------38// ---------------------------------------------------------------------------
25// Shared types39// Shared types
26// ---------------------------------------------------------------------------40// ---------------------------------------------------------------------------
@@ -73,9 +87,10 @@ impl<'de> Deserialize<'de> for LocalBlockHash {
73)]87)]
74#[serde(rename_all = "snake_case")]88#[serde(rename_all = "snake_case")]
75pub enum StorageMedium {89pub enum StorageMedium {
76- /// GPU/NPU HBM — from inference engine workers.90+ /// NPU HBM — from inference engine workers.
91+ /// Wire aliases `gpu` / `xpu` / `hbm` / `device` still parse to this variant.
77 #[default]92 #[default]
78- Xpu,93+ Npu,
79 /// Host DDR / CPU pinned memory — from Mooncake master (MEMORY replica).94 /// Host DDR / CPU pinned memory — from Mooncake master (MEMORY replica).
80 Cpu,95 Cpu,
81 /// SSD / DFS / NVMe disk — from Mooncake master (DISK replica).96 /// SSD / DFS / NVMe disk — from Mooncake master (DISK replica).
@@ -87,7 +102,9 @@ pub enum StorageMedium {
87impl StorageMedium {102impl StorageMedium {
88 pub fn parse(s: &str) -> Self {103 pub fn parse(s: &str) -> Self {
89 match s {104 match s {
90- "xpu" | "XPU" | "hbm" | "HBM" | "device" | "DEVICE" => Self::Xpu,105+ "npu" | "NPU" | "gpu" | "GPU" | "xpu" | "XPU" | "hbm" | "HBM" | "device" | "DEVICE" => {
106+ Self::Npu
107+ }
91 "cpu" | "CPU" | "cpu_pinned" | "CPU_PINNED" | "host" | "HOST" | "memory" | "MEMORY" => {108 "cpu" | "CPU" | "cpu_pinned" | "CPU_PINNED" | "host" | "HOST" | "memory" | "MEMORY" => {
92 Self::Cpu109 Self::Cpu
93 }110 }
@@ -100,36 +117,16 @@ impl StorageMedium {
100 117 
101 pub fn as_str(&self) -> &'static str {118 pub fn as_str(&self) -> &'static str {
102 match self {119 match self {
103- Self::Xpu => "XPU",120+ Self::Npu => "NPU",
104 Self::Cpu => "CPU",121 Self::Cpu => "CPU",
105 Self::Disk => "DISK",122 Self::Disk => "DISK",
106 Self::Unknown => "UNKNOWN",123 Self::Unknown => "UNKNOWN",
107 }124 }
108 }125 }
109-}
110 126 
111-// ---------------------------------------------------------------------------127+ /// Whether `s` names the device-HBM tier (npu / gpu / xpu / hbm / device).
112-// Scoring configuration128+ pub fn is_hbm_key(s: &str) -> bool {
113-// ---------------------------------------------------------------------------129+ matches!(Self::parse(s), Self::Npu)
114- 
115-/// Per-medium block match weights, configurable at startup.
116-#[derive(Debug, Clone)]
117-pub struct ScoringConfig {
118- /// Weight for each HBM (XPU) block matched.
119- pub hbm_weight: u32,
120- /// Weight for each CPU block matched.
121- pub cpu_weight: u32,
122- /// Weight for each disk block matched.
123- pub disk_weight: u32,
124-}
125- 
126-impl Default for ScoringConfig {
127- fn default() -> Self {
128- Self {
129- hbm_weight: 3,
130- cpu_weight: 2,
131- disk_weight: 1,
132- }
133 }130 }
134}131}
135 132 
@@ -150,7 +147,7 @@ pub struct WorkerKey {
150 /// RFC #1527: backend that owns the KV blocks (engine worker, Mooncake daemon, etc.).147 /// RFC #1527: backend that owns the KV blocks (engine worker, Mooncake daemon, etc.).
151 pub backend_id: String,148 pub backend_id: String,
152 pub dp_rank: DpRank,149 pub dp_rank: DpRank,
153- /// RFC #1527: cache medium (xpu, cpu, disk).150+ /// RFC #1527: cache medium (npu, cpu, disk).
154 pub medium: StorageMedium,151 pub medium: StorageMedium,
155}152}
156 153 
@@ -163,7 +160,7 @@ pub struct WorkerKey {
163pub struct RegisterRequest {160pub struct RegisterRequest {
164 pub instance_id: InstanceId,161 pub instance_id: InstanceId,
165 /// Per-medium ZMQ PUB endpoints (new protocol).162 /// Per-medium ZMQ PUB endpoints (new protocol).
166- /// e.g. {"xpu": "tcp://...:5557", "cpu": "tcp://...:5558"}.163+ /// e.g. {"npu": "tcp://...:5557", "cpu": "tcp://...:5558"}.
167 /// Multiple media may share the same endpoint URL; the conductor deduplicates.164 /// Multiple media may share the same endpoint URL; the conductor deduplicates.
168 /// When empty, falls back to the legacy `endpoint` field.165 /// When empty, falls back to the legacy `endpoint` field.
169 #[serde(default)]166 #[serde(default)]
@@ -236,34 +233,27 @@ pub struct QueryByHashRequest {
236// Query response types (matching Python client expectations)233// Query response types (matching Python client expectations)
237//234//
238// Python reads:235// Python reads:
239-// rsp[tenant_id][instance_id]["longest_matched"] (in tokens)236+// rsp[tenant_id][instance_id]["longest_matched"] (in tokens)
240-// rsp[tenant_id][instance_id]["DP"][dp_rank_str] (in tokens)237+// rsp[tenant_id][instance_id]["DP"][dp_rank_str] (DpBlocks obj:
238+// matched_tokens / npu_blocks / cpu_blocks / disk_blocks)
241// ---------------------------------------------------------------------------239// ---------------------------------------------------------------------------
242 240 
243-/// Per-DP weighted score breakdown across storage media.241+/// Per-DP matched block counts across storage media.
244#[derive(Debug, Clone, Serialize, Default)]242#[derive(Debug, Clone, Serialize, Default)]
245-pub struct DpScoring {243+pub struct DpBlocks {
246- /// HBM matched blocks × 3244+ /// Cached prefix length for this DP, in tokens:
247- #[serde(rename = "XPU")]245+ /// farthest absolute coverage end across media × `block_size`.
248- pub xpu_score: u32,246+ ///
249- /// CPU matched blocks × 2247+ /// Per-medium `*_blocks` are segment lengths (HBM from root; CPU/Disk
250- #[serde(rename = "CPU")]248+ /// from their continuation **or** an unconditional root chain — replicas
251- pub cpu_score: u32,249+ /// are reported on every tier). `matched_tokens` must not sum
252- /// Disk matched blocks × 1250+ /// overlapping replicas of the same prefix.
253- #[serde(rename = "DISK")]
254- pub disk_score: u32,
255- /// Total weighted score (xpu_score + cpu_score + disk_score)
256- pub total: u32,
257- /// Cached prefix length for this DP, in tokens (blocks × block_size).
258 pub matched_tokens: u32,251 pub matched_tokens: u32,
259- /// HBM raw block count.252+ /// HBM (NPU) matched block count.
260- #[serde(rename = "XPU_blk")]253+ pub npu_blocks: u32,
261- pub xpu_blocks: u32,254+ /// CPU matched block count.
262- /// CPU raw block count.
263- #[serde(rename = "CPU_blk")]
264 pub cpu_blocks: u32,255 pub cpu_blocks: u32,
265- /// Disk raw block count.256+ /// Disk matched block count.
266- #[serde(rename = "DISK_blk")]
267 pub disk_blocks: u32,257 pub disk_blocks: u32,
268}258}
269 259 
@@ -272,11 +262,9 @@ pub struct DpScoring {
272pub struct InstanceMatchData {262pub struct InstanceMatchData {
273 /// Longest continuous prefix match across all DP ranks, in tokens.263 /// Longest continuous prefix match across all DP ranks, in tokens.
274 pub longest_matched: u32,264 pub longest_matched: u32,
275- /// Per-DP-rank scoring breakdown across media.265+ /// Per-DP-rank matched block counts across media.
276 #[serde(rename = "DP")]266 #[serde(rename = "DP")]
277- pub dp: HashMap<String, DpScoring>,267+ pub dp: HashMap<String, DpBlocks>,
278- /// Sum of all DP total scores for this instance.
279- pub total_score: u32,
280}268}
281 269 
282/// Full query response: { tenant_id: { instance_id: InstanceMatchData } }270/// Full query response: { tenant_id: { instance_id: InstanceMatchData } }
@@ -292,12 +280,12 @@ pub struct QueryResponse {
292 280 
293/// Batch of KV cache events from workers.281/// Batch of KV cache events from workers.
294///282///
295-/// Routing context (`instance_id`, `model_name`, `tenant_id`, `block_size`)283+/// Routing context (`instance_id`, `model_name`, `tenant_id`) identifies the
296-/// identifies the originating worker and the model/tenant scope for indexer284+/// originating worker and the model/tenant scope for indexer lookup. When
297-/// lookup. When `model_name` / `tenant_id` are omitted and the instance is285+/// `model_name` / `tenant_id` are omitted and the instance is already
298-/// already registered, the registered values are used as a fallback.286+/// registered, the registered values are used as a fallback. `block_size`
299-/// `block_size` defaults to 128 when neither the batch nor a prior287+/// is carried for wire compatibility; the HTTP events path does not consume
300-/// registration provides it.288+/// it (query uses the registered value).
301#[derive(Debug, Clone, Deserialize)]289#[derive(Debug, Clone, Deserialize)]
302pub struct KvEventBatch {290pub struct KvEventBatch {
303 /// The worker instance these events originate from.291 /// The worker instance these events originate from.
@@ -308,7 +296,9 @@ pub struct KvEventBatch {
308 /// Tenant id for indexer routing (falls back to registered value if omitted).296 /// Tenant id for indexer routing (falls back to registered value if omitted).
309 #[serde(default)]297 #[serde(default)]
310 pub tenant_id: Option<String>,298 pub tenant_id: Option<String>,
311- /// KV block size in tokens (falls back to registered value, then 128).299+ /// KV block size in tokens. The HTTP events path currently does not
300+ /// consume this field (query uses the registered `block_size`); the
301+ /// serde default of 128 keeps the wire type stable.
312 #[serde(default = "default_block_size")]302 #[serde(default = "default_block_size")]
313 pub block_size: u32,303 pub block_size: u32,
314 #[serde(default)]304 #[serde(default)]
@@ -355,14 +345,24 @@ pub struct KvEventWirePayload {
355 /// Engine-style: blocks with block_hash + tokens_hash.345 /// Engine-style: blocks with block_hash + tokens_hash.
356 pub blocks: Vec<KvCacheStoredBlockData>,346 pub blocks: Vec<KvCacheStoredBlockData>,
357 /// Engine-style: parent sequence hash.347 /// Engine-style: parent sequence hash.
348+ /// Accepts both the RFC #1527 field name `parent_hash` and the vLLM
349+ /// engine field name `parent_block_hash`.
350+ #[serde(alias = "parent_block_hash")]
358 pub parent_hash: Option<i64>,351 pub parent_hash: Option<i64>,
352+ /// Engine-style: raw token ids of the stored chain, used to recompute
353+ /// `tokens_hash` (XXH3) when the event carries no pre-computed blocks.
354+ #[serde(default)]
355+ pub token_ids: Vec<i64>,
356+ /// Engine-style: block size in tokens, must pair with `token_ids`.
357+ #[serde(default)]
358+ pub block_size: Option<u32>,
359 /// RFC #1527: rolling sequence hashes.359 /// RFC #1527: rolling sequence hashes.
360 pub seq_hashes: Vec<u64>,360 pub seq_hashes: Vec<u64>,
361- /// RFC #1527: cache medium (xpu, cpu, disk).361+ /// RFC #1527: cache medium (npu, cpu, disk).
362 pub medium: Option<String>,362 pub medium: Option<String>,
363 /// RFC #1527: backend that owns the blocks.363 /// RFC #1527: backend that owns the blocks.
364 pub backend_id: Option<String>,364 pub backend_id: Option<String>,
365- /// Legacy compat: block_hashes (vLLM/Dynamo alias for seq_hashes).365+ /// Legacy compat: block_hashes (vLLM / engine alias for seq_hashes).
366 #[serde(default)]366 #[serde(default)]
367 pub block_hashes: Vec<u64>,367 pub block_hashes: Vec<u64>,
368 /// Legacy compat: old event type string (e.g. "BlockStored").368 /// Legacy compat: old event type string (e.g. "BlockStored").
@@ -387,6 +387,18 @@ impl KvEventWirePayload {
387 "stored" => {387 "stored" => {
388 let blocks: Vec<KvCacheStoredBlockData> = if !self.blocks.is_empty() {388 let blocks: Vec<KvCacheStoredBlockData> = if !self.blocks.is_empty() {
389 self.blocks.clone()389 self.blocks.clone()
390+ } else if !self.token_ids.is_empty() && self.block_size.is_some_and(|bs| bs > 0) {
391+ // Engine-style event (vLLM map/JSON): recompute the XXH3
392+ // content hash from token_ids, same as the ZMQ path.
393+ let bs = self.block_size.unwrap_or(0);
394+ let computed = compute_block_hash_for_seq(&self.token_ids, bs);
395+ let num = computed.len().min(self.block_hashes.len());
396+ (0..num)
397+ .map(|i| KvCacheStoredBlockData {
398+ block_hash: self.block_hashes[i],
399+ tokens_hash: computed[i].0,
400+ })
401+ .collect()
390 } else {402 } else {
391 // Build engine-style blocks from seq_hashes (Mooncake path: no tokens_hash)403 // Build engine-style blocks from seq_hashes (Mooncake path: no tokens_hash)
392 seq_hashes404 seq_hashes
@@ -397,13 +409,10 @@ impl KvEventWirePayload {
397 })409 })
398 .collect()410 .collect()
399 };411 };
400- let parent_hash = if !self.blocks.is_empty() {
401- self.parent_hash
402- } else {
403- None // Mooncake events have no parent
404- };
405 KvCacheEventData::Stored(KvCacheStoreData {412 KvCacheEventData::Stored(KvCacheStoreData {
406- parent_hash: parent_hash.map(|h| h as u64),413+ // Engine events carry `parent_block_hash` (aliased into
414+ // `parent_hash`); RFC #1527 pool events have no parent.
415+ parent_hash: self.parent_hash.map(|h| h as u64),
407 start_position: None,416 start_position: None,
408 blocks,417 blocks,
409 })418 })
@@ -442,15 +451,31 @@ impl KvEventWirePayload {
442 (data, self.medium.clone(), self.backend_id.clone())451 (data, self.medium.clone(), self.backend_id.clone())
443 }452 }
444 453 
454+ /// Resolve the event type to one of `stored` / `removed` / `cleared`.
455+ ///
456+ /// Both the canonical wire names ("stored", "removed", "cleared") and the
457+ /// engine class names ("BlockStored", "BlockRemoved", "AllBlocksCleared")
458+ /// are recognized via keyword matching, so a vLLM-style event pushed over
459+ /// HTTP /events is never misclassified (a `BlockStored` must not fall into
460+ /// the `Removed` fallback branch).
445 fn resolve_event_type(&self) -> &str {461 fn resolve_event_type(&self) -> &str {
446- if !self.event_type.is_empty() {462+ let raw = if !self.event_type.is_empty() {
447- return &self.event_type;463+ self.event_type.as_str()
448- }464+ } else {
449- match &self.legacy_type {465+ match &self.legacy_type {
450- Some(t) if t.contains("Removed") || t.contains("removed") => "removed",466+ Some(t) => t.as_str(),
451- Some(t) if t.contains("Stored") || t.contains("stored") => "stored",467+ None => return "unknown",
452- Some(t) if t.contains("Cleared") || t.contains("cleared") => "cleared",468+ }
453- _ => "unknown",469+ };
470+ let lower = raw.to_ascii_lowercase();
471+ if lower.contains("removed") {
472+ "removed"
473+ } else if lower.contains("stored") {
474+ "stored"
475+ } else if lower.contains("cleared") {
476+ "cleared"
477+ } else {
478+ "unknown"
454 }479 }
455 }480 }
456 481 
@@ -503,42 +528,42 @@ pub struct KvCacheStoredBlockData {
503// Internal types528// Internal types
504// ---------------------------------------------------------------------------529// ---------------------------------------------------------------------------
505 530 
506-/// Overlap scores result from a radix-tree lookup.531+/// Overlap match result from a radix-tree lookup.
507#[derive(Debug, Clone, Default)]532#[derive(Debug, Clone, Default)]
508-pub struct OverlapScores {533+pub struct OverlapBlocks {
509 /// worker -> matched block count534 /// worker -> matched block count
510- pub scores: FxHashMap<WorkerKey, u32>,535+ pub blocks: FxHashMap<WorkerKey, u32>,
511}536}
512 537 
513-impl OverlapScores {538+impl OverlapBlocks {
514 pub fn is_empty(&self) -> bool {539 pub fn is_empty(&self) -> bool {
515- self.scores.is_empty()540+ self.blocks.is_empty()
516 }541 }
517 542 
518- /// Add points for a worker (HBM depth × 3, CPU block × 2, disk block × 1).543+ /// Add matched block count for a worker.
519 #[inline]544 #[inline]
520- pub fn add_score(&mut self, worker: WorkerKey, points: u32) {545+ pub fn add_blocks(&mut self, worker: WorkerKey, n: u32) {
521- self.scores546+ self.blocks
522 .entry(worker)547 .entry(worker)
523- .and_modify(|s| *s += points)548+ .and_modify(|s| *s += n)
524- .or_insert(points);549+ .or_insert(n);
525 }550 }
526 551 
527- /// Legacy max-based update, used internally by HBM tree traversal.552+ /// Max-based update, used internally by HBM tree traversal.
528 #[inline]553 #[inline]
529- pub fn update_score(&mut self, worker: WorkerKey, depth: u32) {554+ pub fn update_blocks(&mut self, worker: WorkerKey, depth: u32) {
530- self.scores555+ self.blocks
531 .entry(worker)556 .entry(worker)
532 .and_modify(|s| *s = (*s).max(depth))557 .and_modify(|s| *s = (*s).max(depth))
533 .or_insert(depth);558 .or_insert(depth);
534 }559 }
535 560 
536- /// Merge scores from another `OverlapScores` into this one.561+ /// Merge matched blocks from another `OverlapBlocks` into this one,
537- /// Used by parallel flat lookup to combine per-thread results.562+ /// keeping the maximum per worker.
538 #[inline]563 #[inline]
539- pub fn merge(&mut self, other: OverlapScores) {564+ pub fn merge(&mut self, other: OverlapBlocks) {
540- for (worker, score) in other.scores {565+ for (worker, n) in other.blocks {
541- self.add_score(worker, score);566+ self.add_blocks(worker, n);
542 }567 }
543 }568 }
544}569}
@@ -553,11 +578,18 @@ mod tests {
553 // ── StorageMedium ─────────────────────────────────────────────────578 // ── StorageMedium ─────────────────────────────────────────────────
554 579 
555 #[test]580 #[test]
556- fn test_storage_medium_from_str_xpu() {581+ fn test_storage_medium_from_str_npu() {
557- assert_eq!(StorageMedium::parse("xpu"), StorageMedium::Xpu);582+ assert_eq!(StorageMedium::parse("npu"), StorageMedium::Npu);
558- assert_eq!(StorageMedium::parse("XPU"), StorageMedium::Xpu);583+ assert_eq!(StorageMedium::parse("NPU"), StorageMedium::Npu);
559- assert_eq!(StorageMedium::parse("hbm"), StorageMedium::Xpu);584+ assert_eq!(StorageMedium::parse("gpu"), StorageMedium::Npu);
560- assert_eq!(StorageMedium::parse("device"), StorageMedium::Xpu);585+ assert_eq!(StorageMedium::parse("GPU"), StorageMedium::Npu);
586+ assert_eq!(StorageMedium::parse("xpu"), StorageMedium::Npu);
587+ assert_eq!(StorageMedium::parse("XPU"), StorageMedium::Npu);
588+ assert_eq!(StorageMedium::parse("hbm"), StorageMedium::Npu);
589+ assert_eq!(StorageMedium::parse("device"), StorageMedium::Npu);
590+ assert!(StorageMedium::is_hbm_key("npu"));
591+ assert!(StorageMedium::is_hbm_key("gpu"));
592+ assert!(!StorageMedium::is_hbm_key("cpu"));
561 }593 }
562 594 
563 #[test]595 #[test]
@@ -573,13 +605,13 @@ mod tests {
573 }605 }
574 606 
575 #[test]607 #[test]
576- fn test_storage_medium_default_is_xpu() {608+ fn test_storage_medium_default_is_npu() {
577- assert_eq!(StorageMedium::default(), StorageMedium::Xpu);609+ assert_eq!(StorageMedium::default(), StorageMedium::Npu);
578 }610 }
579 611 
580 #[test]612 #[test]
581 fn test_storage_medium_as_str() {613 fn test_storage_medium_as_str() {
582- assert_eq!(StorageMedium::Xpu.as_str(), "XPU");614+ assert_eq!(StorageMedium::Npu.as_str(), "NPU");
583 assert_eq!(StorageMedium::Cpu.as_str(), "CPU");615 assert_eq!(StorageMedium::Cpu.as_str(), "CPU");
584 assert_eq!(StorageMedium::Disk.as_str(), "DISK");616 assert_eq!(StorageMedium::Disk.as_str(), "DISK");
585 assert_eq!(StorageMedium::Unknown.as_str(), "UNKNOWN");617 assert_eq!(StorageMedium::Unknown.as_str(), "UNKNOWN");
@@ -593,12 +625,12 @@ mod tests {
593 instance_id: "inst-1".into(),625 instance_id: "inst-1".into(),
594 backend_id: "backend-a".into(),626 backend_id: "backend-a".into(),
595 dp_rank: 2,627 dp_rank: 2,
596- medium: StorageMedium::Xpu,628+ medium: StorageMedium::Npu,
597 };629 };
598 assert_eq!(wk.instance_id, "inst-1");630 assert_eq!(wk.instance_id, "inst-1");
599 assert_eq!(wk.backend_id, "backend-a");631 assert_eq!(wk.backend_id, "backend-a");
600 assert_eq!(wk.dp_rank, 2);632 assert_eq!(wk.dp_rank, 2);
601- assert_eq!(wk.medium, StorageMedium::Xpu);633+ assert_eq!(wk.medium, StorageMedium::Npu);
602 }634 }
603 635 
604 #[test]636 #[test]
@@ -624,7 +656,7 @@ mod tests {
624 instance_id: "i1".into(),656 instance_id: "i1".into(),
625 backend_id: "b1".into(),657 backend_id: "b1".into(),
626 dp_rank: 0,658 dp_rank: 0,
627- medium: StorageMedium::Xpu,659+ medium: StorageMedium::Npu,
628 };660 };
629 let b = WorkerKey {661 let b = WorkerKey {
630 instance_id: "i1".into(),662 instance_id: "i1".into(),
@@ -639,50 +671,41 @@ mod tests {
639 671 
640 #[test]672 #[test]
641 fn test_instance_match_data_serialization() {673 fn test_instance_match_data_serialization() {
642- let mut imd = InstanceMatchData::default();674+ let mut imd = InstanceMatchData {
643- imd.longest_matched = 256;675+ longest_matched: 256,
676+ ..Default::default()
677+ };
644 imd.dp.insert(678 imd.dp.insert(
645 "0".into(),679 "0".into(),
646- DpScoring {680+ DpBlocks {
647- xpu_score: 6,
648- cpu_score: 0,
649- disk_score: 0,
650- total: 6,
651 matched_tokens: 768,681 matched_tokens: 768,
652- xpu_blocks: 6,682+ npu_blocks: 6,
653 cpu_blocks: 0,683 cpu_blocks: 0,
654 disk_blocks: 0,684 disk_blocks: 0,
655 },685 },
656 );686 );
657 imd.dp.insert(687 imd.dp.insert(
658 "1".into(),688 "1".into(),
659- DpScoring {689+ DpBlocks {
660- xpu_score: 0,
661- cpu_score: 8,
662- disk_score: 0,
663- total: 8,
664 matched_tokens: 1024,690 matched_tokens: 1024,
665- xpu_blocks: 0,691+ npu_blocks: 0,
666 cpu_blocks: 4,692 cpu_blocks: 4,
667 disk_blocks: 0,693 disk_blocks: 0,
668 },694 },
669 );695 );
670- imd.total_score = 14;
671 696 
672 let json = serde_json::to_string(&imd).unwrap();697 let json = serde_json::to_string(&imd).unwrap();
673 let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();698 let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
674 699 
675 assert_eq!(parsed["longest_matched"], 256);700 assert_eq!(parsed["longest_matched"], 256);
676- assert_eq!(parsed["total_score"], 14);701+ assert!(parsed.get("total_score").is_none());
677- assert_eq!(parsed["DP"]["0"]["XPU"], 6);702+ assert!(parsed["DP"]["0"].get("XPU").is_none());
678- assert_eq!(parsed["DP"]["0"]["total"], 6);703+ assert!(parsed["DP"]["0"].get("total").is_none());
679 assert_eq!(parsed["DP"]["0"]["matched_tokens"], 768);704 assert_eq!(parsed["DP"]["0"]["matched_tokens"], 768);
680- assert_eq!(parsed["DP"]["0"]["XPU_blk"], 6);705+ assert_eq!(parsed["DP"]["0"]["npu_blocks"], 6);
681- assert!(parsed["DP"]["0"]["CPU_blk"].as_u64().unwrap() == 0);706+ assert_eq!(parsed["DP"]["0"]["cpu_blocks"], 0);
682- assert_eq!(parsed["DP"]["1"]["CPU"], 8);
683- assert_eq!(parsed["DP"]["1"]["total"], 8);
684 assert_eq!(parsed["DP"]["1"]["matched_tokens"], 1024);707 assert_eq!(parsed["DP"]["1"]["matched_tokens"], 1024);
685- assert_eq!(parsed["DP"]["1"]["CPU_blk"], 4);708+ assert_eq!(parsed["DP"]["1"]["cpu_blocks"], 4);
686 }709 }
687 710 
688 // ── KvEventWirePayload normalization ────────────────────────────────711 // ── KvEventWirePayload normalization ────────────────────────────────
@@ -758,4 +781,72 @@ mod tests {
758 let (data, _, _) = payload.normalize();781 let (data, _, _) = payload.normalize();
759 assert!(matches!(data, KvCacheEventData::Removed { .. }));782 assert!(matches!(data, KvCacheEventData::Removed { .. }));
760 }783 }
784+ 
785+ // ── Engine class-name type resolution (vLLM map/JSON events) ────────
786+ 
787+ #[test]
788+ fn test_normalize_engine_block_stored_is_stored_not_removed() {
789+ // vLLM engine events pushed via HTTP /events carry `type: "BlockStored"`.
790+ // They must normalize to Stored — the historical fallback classified
791+ // them as Removed, which would delete index entries.
792+ let payload = KvEventWirePayload {
793+ event_type: "BlockStored".into(),
794+ block_hashes: vec![100, 200],
795+ parent_hash: Some(50),
796+ token_ids: vec![1, 2, 3, 4, 5, 6, 7, 8],
797+ block_size: Some(4),
798+ ..Default::default()
799+ };
800+ let (data, _, _) = payload.normalize();
801+ let KvCacheEventData::Stored(store) = &data else {
802+ panic!("BlockStored must normalize to Stored, got {data:?}");
803+ };
804+ assert_eq!(store.parent_hash, Some(50));
805+ assert_eq!(store.blocks.len(), 2);
806+ // tokens_hash is the XXH3 content hash recomputed from token_ids,
807+ // never the sequence hash.
808+ let computed = compute_block_hash_for_seq(&[1, 2, 3, 4, 5, 6, 7, 8], 4);
809+ assert_eq!(store.blocks[0].tokens_hash, computed[0].0);
810+ assert_eq!(store.blocks[1].tokens_hash, computed[1].0);
811+ assert_ne!(store.blocks[0].tokens_hash, 100);
812+ }
813+ 
814+ #[test]
815+ fn test_normalize_engine_block_removed() {
816+ let payload = KvEventWirePayload {
817+ event_type: "BlockRemoved".into(),
818+ block_hashes: vec![111],
819+ ..Default::default()
820+ };
821+ let (data, _, _) = payload.normalize();
822+ assert!(matches!(data, KvCacheEventData::Removed { .. }));
823+ }
824+ 
825+ #[test]
826+ fn test_normalize_engine_all_blocks_cleared() {
827+ let payload = KvEventWirePayload {
828+ event_type: "AllBlocksCleared".into(),
829+ ..Default::default()
830+ };
831+ let (data, _, _) = payload.normalize();
832+ assert!(matches!(data, KvCacheEventData::Cleared));
833+ }
834+ 
835+ #[test]
836+ fn test_normalize_engine_parent_block_hash_alias() {
837+ // The vLLM field name `parent_block_hash` maps onto `parent_hash`.
838+ let payload = KvEventWirePayload {
839+ event_type: "BlockStored".into(),
840+ block_hashes: vec![100],
841+ parent_hash: Some(999),
842+ token_ids: vec![1, 2, 3, 4],
843+ block_size: Some(4),
844+ ..Default::default()
845+ };
846+ let (data, _, _) = payload.normalize();
847+ let KvCacheEventData::Stored(store) = &data else {
848+ panic!("expected Stored, got {data:?}");
849+ };
850+ assert_eq!(store.parent_hash, Some(999));
851+ }
761}852}
Mmotor/kv_conductor/src/registry.rs+15-15
@@ -38,7 +38,7 @@ fn extract_ip_from_endpoint(endpoint: &str) -> Option<String> {
38 Some(without_prefix[..colon_pos].to_string())38 Some(without_prefix[..colon_pos].to_string())
39}39}
40 40 
41-/// Add `(instance_id, dp_rank)` entries to `hbm_ip_index` for each XPU endpoint.41+/// Add `(instance_id, dp_rank)` entries to `hbm_ip_index` for each NPU endpoint.
42fn add_hbm_ip_index_entries(42fn add_hbm_ip_index_entries(
43 ip_index: &HbmIpIndex,43 ip_index: &HbmIpIndex,
44 medium_endpoints: &HashMap<String, String>,44 medium_endpoints: &HashMap<String, String>,
@@ -46,7 +46,7 @@ fn add_hbm_ip_index_entries(
46 dp_rank: u32,46 dp_rank: u32,
47) {47) {
48 for (medium_str, ep_url) in medium_endpoints {48 for (medium_str, ep_url) in medium_endpoints {
49- if !medium_str.eq_ignore_ascii_case("xpu") {49+ if !StorageMedium::is_hbm_key(medium_str) {
50 continue;50 continue;
51 }51 }
52 let Some(ref ip) = extract_ip_from_endpoint(ep_url) else {52 let Some(ref ip) = extract_ip_from_endpoint(ep_url) else {
@@ -61,7 +61,7 @@ fn add_hbm_ip_index_entries(
61 }61 }
62}62}
63 63 
64-/// For each XPU endpoint in `medium_endpoints`, remove the (instance_id, dp_rank)64+/// For each NPU endpoint in `medium_endpoints`, remove the (instance_id, dp_rank)
65/// entry from `hbm_ip_index`. If a given IP has no remaining DPs, the IP key is65/// entry from `hbm_ip_index`. If a given IP has no remaining DPs, the IP key is
66/// dropped entirely.66/// dropped entirely.
67fn remove_hbm_ip_index_entries(67fn remove_hbm_ip_index_entries(
@@ -71,7 +71,7 @@ fn remove_hbm_ip_index_entries(
71 dp_rank: u32,71 dp_rank: u32,
72) {72) {
73 for (medium_str, ep_url) in medium_endpoints {73 for (medium_str, ep_url) in medium_endpoints {
74- if !medium_str.eq_ignore_ascii_case("xpu") {74+ if !StorageMedium::is_hbm_key(medium_str) {
75 continue;75 continue;
76 }76 }
77 let Some(ref ip) = extract_ip_from_endpoint(ep_url) else {77 let Some(ref ip) = extract_ip_from_endpoint(ep_url) else {
@@ -91,7 +91,7 @@ fn remove_hbm_ip_index_entries(
91/// Information about a registered endpoint for a worker.91/// Information about a registered endpoint for a worker.
92#[derive(Debug, Clone, Serialize)]92#[derive(Debug, Clone, Serialize)]
93pub struct EndpointInfo {93pub struct EndpointInfo {
94- /// Per-medium ZMQ PUB endpoints, e.g. {"xpu": "tcp://...:5557", "cpu": "tcp://...:5558"}.94+ /// Per-medium ZMQ PUB endpoints, e.g. {"npu": "tcp://...:5557", "cpu": "tcp://...:5558"}.
95 pub medium_endpoints: HashMap<String, String>,95 pub medium_endpoints: HashMap<String, String>,
96 pub engine_type: String,96 pub engine_type: String,
97 pub dp_rank: DpRank,97 pub dp_rank: DpRank,
@@ -119,8 +119,8 @@ pub struct WorkerRegistry {
119 instances: tokio::sync::RwLock<HashMap<InstanceId, WorkerEntry>>,119 instances: tokio::sync::RwLock<HashMap<InstanceId, WorkerEntry>>,
120 /// The shared indexer for all models/tenants120 /// The shared indexer for all models/tenants
121 indexer: Arc<Indexer>,121 indexer: Arc<Indexer>,
122- /// Count of active replay sessions. Queries are rejected while > 0122+ /// Count of active replay sessions. Tracks how many replays are running
123- /// to avoid returning incomplete results during prefix-tree rebuild.123+ /// (for diagnostics); replay itself is offloaded to `spawn_blocking`.
124 /// Wrapped in Arc so it can be shared with spawn_blocking tasks.124 /// Wrapped in Arc so it can be shared with spawn_blocking tasks.
125 replay_in_progress: Arc<std::sync::atomic::AtomicU64>,125 replay_in_progress: Arc<std::sync::atomic::AtomicU64>,
126 /// Active ZMQ subscribers, keyed by (instance_id, dp_rank, endpoint_url).126 /// Active ZMQ subscribers, keyed by (instance_id, dp_rank, endpoint_url).
@@ -164,12 +164,12 @@ impl WorkerRegistry {
164 ));164 ));
165 }165 }
166 let mut map: HashMap<String, Vec<StorageMedium>> = HashMap::new();166 let mut map: HashMap<String, Vec<StorageMedium>> = HashMap::new();
167- // Mooncake master publishes CPU and DISK events (never XPU) on one167+ // Mooncake master publishes CPU and DISK events (never NPU) on one
168- // port. We include XPU here so the indexer is ready for HBM events168+ // port. We include NPU here so the indexer is ready for HBM events
169 // from engine-level publishers that may share the same endpoint.169 // from engine-level publishers that may share the same endpoint.
170 map.insert(170 map.insert(
171 ep.clone(),171 ep.clone(),
172- vec![StorageMedium::Xpu, StorageMedium::Cpu, StorageMedium::Disk],172+ vec![StorageMedium::Npu, StorageMedium::Cpu, StorageMedium::Disk],
173 );173 );
174 return Ok(map);174 return Ok(map);
175 }175 }
@@ -179,10 +179,10 @@ impl WorkerRegistry {
179 Ok(HashMap::new())179 Ok(HashMap::new())
180 }180 }
181 181 
182- pub fn new(scoring: ScoringConfig) -> Self {182+ pub fn new() -> Self {
183 Self {183 Self {
184 instances: tokio::sync::RwLock::new(HashMap::new()),184 instances: tokio::sync::RwLock::new(HashMap::new()),
185- indexer: Arc::new(Indexer::new(scoring)),185+ indexer: Arc::new(Indexer::new()),
186 replay_in_progress: Arc::new(std::sync::atomic::AtomicU64::new(0)),186 replay_in_progress: Arc::new(std::sync::atomic::AtomicU64::new(0)),
187 zmq_subscribers: tokio::sync::RwLock::new(HashMap::new()),187 zmq_subscribers: tokio::sync::RwLock::new(HashMap::new()),
188 hbm_ip_index: Arc::new(ParkingRwLock::new(HashMap::new())),188 hbm_ip_index: Arc::new(ParkingRwLock::new(HashMap::new())),
@@ -449,7 +449,7 @@ impl WorkerRegistry {
449 449 
450 entry.endpoints.remove(&req.dp_rank);450 entry.endpoints.remove(&req.dp_rank);
451 451 
452- // Remove from indexer tree across all storage media (XPU/CPU/DISK)452+ // Remove from indexer tree across all storage media (NPU/CPU/DISK)
453 // Use the *registered* model/tenant (not the request's) to ensure453 // Use the *registered* model/tenant (not the request's) to ensure
454 // cleanup targets the correct indexer entry.454 // cleanup targets the correct indexer entry.
455 let indexer_entry = self.indexer.get(&entry.model_name, &entry.tenant_id);455 let indexer_entry = self.indexer.get(&entry.model_name, &entry.tenant_id);
@@ -548,7 +548,7 @@ impl WorkerRegistry {
548 let medium = medium_str548 let medium = medium_str
549 .as_deref()549 .as_deref()
550 .map(StorageMedium::parse)550 .map(StorageMedium::parse)
551- .unwrap_or(StorageMedium::Xpu);551+ .unwrap_or(StorageMedium::Npu);
552 552 
553 let backend_id = backend_id_opt.as_deref().unwrap_or(instance_id).to_string();553 let backend_id = backend_id_opt.as_deref().unwrap_or(instance_id).to_string();
554 554 
@@ -594,7 +594,7 @@ impl WorkerRegistry {
594 594 
595impl Default for WorkerRegistry {595impl Default for WorkerRegistry {
596 fn default() -> Self {596 fn default() -> Self {
597- Self::new(ScoringConfig::default())597+ Self::new()
598 }598 }
599}599}
600 600 
Mmotor/kv_conductor/src/server.rs+10-9
@@ -13,8 +13,8 @@
13//! Provides the following endpoints:13//! Provides the following endpoints:
14//! - `POST /register` — Register a worker instance14//! - `POST /register` — Register a worker instance
15//! - `POST /unregister` — Unregister a worker instance15//! - `POST /unregister` — Unregister a worker instance
16-//! - `POST /query` — Query KV cache overlap scores by token IDs16+//! - `POST /query` — Query KV cache matched blocks by token IDs
17-//! - `POST /query_by_hash` — Query KV cache overlap scores by pre-computed hashes17+//! - `POST /query_by_hash` — Query KV cache matched blocks by pre-computed hashes
18//! - `POST /events` — Ingest KV cache events from workers18//! - `POST /events` — Ingest KV cache events from workers
19//! - `GET /health` — Health check19//! - `GET /health` — Health check
20//! - `GET /workers` — List registered workers (debug)20//! - `GET /workers` — List registered workers (debug)
@@ -30,7 +30,7 @@ use axum::{
30use tower_http::cors::{Any, CorsLayer};30use tower_http::cors::{Any, CorsLayer};
31use tower_http::trace::TraceLayer;31use tower_http::trace::TraceLayer;
32 32 
33-/// Maximum request body size (16 MB). Large queries (402400+ token IDs)33+/// Maximum request body size (64 MB). Large queries (402400+ token IDs)
34/// exceed axum's default 2 MB limit.34/// exceed axum's default 2 MB limit.
35const MAX_BODY_BYTES: usize = 64 * 1024 * 1024;35const MAX_BODY_BYTES: usize = 64 * 1024 * 1024;
36 36 
@@ -42,8 +42,6 @@ use crate::registry::WorkerRegistry;
42#[derive(Clone)]42#[derive(Clone)]
43pub struct AppState {43pub struct AppState {
44 pub registry: Arc<WorkerRegistry>,44 pub registry: Arc<WorkerRegistry>,
45- /// Per-medium block scoring weights.
46- pub scoring: ScoringConfig,
47}45}
48 46 
49/// Create the axum Router with all endpoints.47/// Create the axum Router with all endpoints.
@@ -62,8 +60,8 @@ pub fn create_router(state: AppState) -> Router {
62 .route("/health", get(health_handler))60 .route("/health", get(health_handler))
63 .route("/workers", get(workers_handler));61 .route("/workers", get(workers_handler));
64 62 
65- // Raise body limit for query/events endpoints — DeepSeek V4 queries63+ // Raise the global body limit — DeepSeek V4 queries carry 400K+ token
66- // carry 400K+ token IDs (~2.4 MB JSON body).64+ // IDs (~2.4 MB JSON body), beyond axum's 2 MB default.
67 router = router.layer(axum::extract::DefaultBodyLimit::max(MAX_BODY_BYTES));65 router = router.layer(axum::extract::DefaultBodyLimit::max(MAX_BODY_BYTES));
68 66 
69 router67 router
@@ -137,7 +135,9 @@ async fn unregister_handler(
137///135///
138/// Request body: `{ "model": "...", "block_size": 128, "token_ids": [...], "tenant_id": "default" }`136/// Request body: `{ "model": "...", "block_size": 128, "token_ids": [...], "tenant_id": "default" }`
139///137///
140-/// Response: `{ "<tenant_id>": { "<instance_id>": { "longest_matched": N, "DP": { "<rank>": N } } } }`138+/// Response: `{ "<tenant_id>": { "<instance_id>": { "longest_matched": N,
139+/// "DP": { "<rank>": { "matched_tokens": N, "npu_blocks": N,
140+/// "cpu_blocks": N, "disk_blocks": N } } } } }`
141async fn query_handler(141async fn query_handler(
142 State(state): State<AppState>,142 State(state): State<AppState>,
143 Json(req): Json<QueryRequest>,143 Json(req): Json<QueryRequest>,
@@ -286,7 +286,8 @@ async fn events_handler(
286 i = j;286 i = j;
287 }287 }
288 288 
289- // Handle shutdown flag: unregister the instance if it's shutting down289+ // Handle shutdown flag: the instance reports it is shutting down. Full
290+ // cleanup is done by an explicit /unregister call; here we just log.
290 if batch.shutdown {291 if batch.shutdown {
291 tracing::info!(292 tracing::info!(
292 instance_id = %batch.instance_id,293 instance_id = %batch.instance_id,
Mmotor/kv_conductor/src/zmq_subscriber.rs+2-2
@@ -419,8 +419,8 @@ fn zmq_errno_reasonable(e: &zmq::Error) -> bool {
419/// 3. End-of-stream: seq == 0xFFFFFFFFFFFFFFFF (-1 as signed i64)419/// 3. End-of-stream: seq == 0xFFFFFFFFFFFFFFFF (-1 as signed i64)
420///420///
421/// Called during `/register` when the registration payload includes421/// Called during `/register` when the registration payload includes
422-/// a `replay_endpoint` field. Runs synchronously in the registration422+/// a `replay_endpoint` field. Invoked via `spawn_blocking` from the
423-/// handler (blocking).423+/// registration handler (offloaded to a dedicated thread).
424#[allow(clippy::too_many_arguments)]424#[allow(clippy::too_many_arguments)]
425pub fn replay_events(425pub fn replay_events(
426 replay_endpoint: &str,426 replay_endpoint: &str,
Mmotor/kv_conductor/tests/integration_test.rs+18-20
@@ -16,15 +16,13 @@ use reqwest::Client;
16use serde_json::{json, Value};16use serde_json::{json, Value};
17use tokio::net::TcpListener;17use tokio::net::TcpListener;
18 18 
19-use kv_conductor::protocols::ScoringConfig;
20use kv_conductor::registry::WorkerRegistry;19use kv_conductor::registry::WorkerRegistry;
21use kv_conductor::server::{create_router, AppState};20use kv_conductor::server::{create_router, AppState};
22 21 
23/// Start a test server on a random port, returning the base URL.22/// Start a test server on a random port, returning the base URL.
24async fn start_test_server() -> (String, tokio::task::JoinHandle<()>) {23async fn start_test_server() -> (String, tokio::task::JoinHandle<()>) {
25- let scoring = ScoringConfig::default();24+ let registry = Arc::new(WorkerRegistry::new());
26- let registry = Arc::new(WorkerRegistry::new(scoring.clone()));25+ let state = AppState { registry };
27- let state = AppState { registry, scoring };
28 let router = create_router(state);26 let router = create_router(state);
29 27 
30 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();28 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -65,7 +63,7 @@ async fn test_register_and_query() {
65 let register_data = json!({63 let register_data = json!({
66 "instance_id": "vllm-prefill-42",64 "instance_id": "vllm-prefill-42",
67 "medium_endpoints": {65 "medium_endpoints": {
68- "xpu": "tcp://10.0.0.1:50090",66+ "npu": "tcp://10.0.0.1:50090",
69 "cpu": "tcp://10.0.0.1:50090",67 "cpu": "tcp://10.0.0.1:50090",
70 "disk": "tcp://10.0.0.1:50090"68 "disk": "tcp://10.0.0.1:50090"
71 },69 },
@@ -120,7 +118,7 @@ async fn test_query_after_kv_events() {
120 .json(&json!({118 .json(&json!({
121 "instance_id": format!("vllm-prefill-{}", i),119 "instance_id": format!("vllm-prefill-{}", i),
122 "medium_endpoints": {120 "medium_endpoints": {
123- "xpu": ep,121+ "npu": ep,
124 "cpu": ep,122 "cpu": ep,
125 "disk": ep123 "disk": ep
126 },124 },
@@ -200,7 +198,7 @@ async fn test_unregister() {
200 .json(&json!({198 .json(&json!({
201 "instance_id": "vllm-prefill-99",199 "instance_id": "vllm-prefill-99",
202 "medium_endpoints": {200 "medium_endpoints": {
203- "xpu": "tcp://10.0.0.1:50090",201+ "npu": "tcp://10.0.0.1:50090",
204 "cpu": "tcp://10.0.0.1:50090",202 "cpu": "tcp://10.0.0.1:50090",
205 "disk": "tcp://10.0.0.1:50090"203 "disk": "tcp://10.0.0.1:50090"
206 },204 },
@@ -260,7 +258,7 @@ async fn test_duplicate_registration() {
260 let reg = json!({258 let reg = json!({
261 "instance_id": "dup-test",259 "instance_id": "dup-test",
262 "medium_endpoints": {260 "medium_endpoints": {
263- "xpu": "tcp://10.0.0.1:50090",261+ "npu": "tcp://10.0.0.1:50090",
264 "cpu": "tcp://10.0.0.1:50090",262 "cpu": "tcp://10.0.0.1:50090",
265 "disk": "tcp://10.0.0.1:50090"263 "disk": "tcp://10.0.0.1:50090"
266 },264 },
@@ -332,12 +330,12 @@ async fn test_mooncake_hbm_plus_pool_registration() {
332 let (base_url, _handle) = start_test_server().await;330 let (base_url, _handle) = start_test_server().await;
333 let client = Client::new();331 let client = Client::new();
334 332 
335- // Register HBM endpoint (XPU only, store_backend=Mooncake)333+ // Register HBM endpoint (NPU only, store_backend=Mooncake)
336 let resp = client334 let resp = client
337 .post(format!("{}/register", base_url))335 .post(format!("{}/register", base_url))
338 .json(&json!({336 .json(&json!({
339 "instance_id": "mooncake-prefill-0",337 "instance_id": "mooncake-prefill-0",
340- "medium_endpoints": {"xpu": "tcp://10.0.0.1:50090"},338+ "medium_endpoints": {"npu": "tcp://10.0.0.1:50090"},
341 "type": "vLLM",339 "type": "vLLM",
342 "store_backend": "Mooncake",340 "store_backend": "Mooncake",
343 "modelname": "mooncake-model",341 "modelname": "mooncake-model",
@@ -395,7 +393,7 @@ async fn test_memcache_hbm_plus_pool_registration() {
395 .post(format!("{}/register", base_url))393 .post(format!("{}/register", base_url))
396 .json(&json!({394 .json(&json!({
397 "instance_id": "memcache-prefill-0",395 "instance_id": "memcache-prefill-0",
398- "medium_endpoints": {"xpu": "tcp://10.0.1.1:50090"},396+ "medium_endpoints": {"npu": "tcp://10.0.1.1:50090"},
399 "type": "vLLM",397 "type": "vLLM",
400 "store_backend": "Memcache",398 "store_backend": "Memcache",
401 "modelname": "memcache-model",399 "modelname": "memcache-model",
@@ -440,13 +438,13 @@ async fn test_yuanrong_multi_port_registration() {
440 let (base_url, _handle) = start_test_server().await;438 let (base_url, _handle) = start_test_server().await;
441 let client = Client::new();439 let client = Client::new();
442 440 
443- // YuanRong: cpu + disk share one port, xpu on another441+ // YuanRong: cpu + disk share one port, npu on another
444 let resp = client442 let resp = client
445 .post(format!("{}/register", base_url))443 .post(format!("{}/register", base_url))
446 .json(&json!({444 .json(&json!({
447 "instance_id": "yr-node-0",445 "instance_id": "yr-node-0",
448 "medium_endpoints": {446 "medium_endpoints": {
449- "xpu": "tcp://10.0.2.1:15557",447+ "npu": "tcp://10.0.2.1:15557",
450 "cpu": "tcp://10.0.2.1:15558",448 "cpu": "tcp://10.0.2.1:15558",
451 "disk": "tcp://10.0.2.1:15558"449 "disk": "tcp://10.0.2.1:15558"
452 },450 },
@@ -470,7 +468,7 @@ async fn test_yuanrong_multi_port_registration() {
470 let body: Value = resp.json().await.unwrap();468 let body: Value = resp.json().await.unwrap();
471 let w = &body["workers"].as_array().unwrap()[0];469 let w = &body["workers"].as_array().unwrap()[0];
472 let meps = &w["endpoints"]["0"]["medium_endpoints"];470 let meps = &w["endpoints"]["0"]["medium_endpoints"];
473- assert_eq!(meps["xpu"], "tcp://10.0.2.1:15557");471+ assert_eq!(meps["npu"], "tcp://10.0.2.1:15557");
474 assert_eq!(meps["cpu"], "tcp://10.0.2.1:15558");472 assert_eq!(meps["cpu"], "tcp://10.0.2.1:15558");
475 assert_eq!(meps["disk"], "tcp://10.0.2.1:15558");473 assert_eq!(meps["disk"], "tcp://10.0.2.1:15558");
476}474}
@@ -484,7 +482,7 @@ async fn test_mooncake_duplicate_hbm_registration() {
484 482 
485 let reg = json!({483 let reg = json!({
486 "instance_id": "mooncake-dup",484 "instance_id": "mooncake-dup",
487- "medium_endpoints": {"xpu": "tcp://10.0.3.1:50090"},485+ "medium_endpoints": {"npu": "tcp://10.0.3.1:50090"},
488 "type": "vLLM",486 "type": "vLLM",
489 "store_backend": "Mooncake",487 "store_backend": "Mooncake",
490 "modelname": "dup-model",488 "modelname": "dup-model",
@@ -557,7 +555,7 @@ async fn test_unregister_mooncake_hbm_removes_worker() {
557 .post(format!("{}/register", base_url))555 .post(format!("{}/register", base_url))
558 .json(&json!({556 .json(&json!({
559 "instance_id": "to-remove",557 "instance_id": "to-remove",
560- "medium_endpoints": {"xpu": "tcp://10.0.5.1:50090"},558+ "medium_endpoints": {"npu": "tcp://10.0.5.1:50090"},
561 "type": "vLLM",559 "type": "vLLM",
562 "store_backend": "Mooncake",560 "store_backend": "Mooncake",
563 "modelname": "rm-model",561 "modelname": "rm-model",
@@ -615,7 +613,7 @@ async fn test_unknown_backend_falls_back_to_yuanrong() {
615 .json(&json!({613 .json(&json!({
616 "instance_id": "unknown-backend",614 "instance_id": "unknown-backend",
617 "medium_endpoints": {615 "medium_endpoints": {
618- "xpu": "tcp://10.0.6.1:15557",616+ "npu": "tcp://10.0.6.1:15557",
619 "cpu": "tcp://10.0.6.1:15558",617 "cpu": "tcp://10.0.6.1:15558",
620 "disk": "tcp://10.0.6.1:15558"618 "disk": "tcp://10.0.6.1:15558"
621 },619 },
@@ -683,7 +681,7 @@ async fn test_reregister_same_backend_preserves_tree() {
683 let reg = json!({681 let reg = json!({
684 "instance_id": "rereg-same",682 "instance_id": "rereg-same",
685 "medium_endpoints": {683 "medium_endpoints": {
686- "xpu": "tcp://10.0.10.1:50090",684+ "npu": "tcp://10.0.10.1:50090",
687 "cpu": "tcp://10.0.10.1:50090",685 "cpu": "tcp://10.0.10.1:50090",
688 "disk": "tcp://10.0.10.1:50090"686 "disk": "tcp://10.0.10.1:50090"
689 },687 },
@@ -727,7 +725,7 @@ async fn test_reregister_same_backend_preserves_tree() {
727 // Re-register with same backend (different endpoint)725 // Re-register with same backend (different endpoint)
728 let mut reg2 = reg.clone();726 let mut reg2 = reg.clone();
729 reg2["medium_endpoints"] = json!({727 reg2["medium_endpoints"] = json!({
730- "xpu": "tcp://10.0.10.2:50090",728+ "npu": "tcp://10.0.10.2:50090",
731 "cpu": "tcp://10.0.10.2:50090",729 "cpu": "tcp://10.0.10.2:50090",
732 "disk": "tcp://10.0.10.2:50090"730 "disk": "tcp://10.0.10.2:50090"
733 });731 });
@@ -775,7 +773,7 @@ async fn test_reregister_different_backend_drops_tree() {
775 let reg = json!({773 let reg = json!({
776 "instance_id": "rereg-diff",774 "instance_id": "rereg-diff",
777 "medium_endpoints": {775 "medium_endpoints": {
778- "xpu": "tcp://10.0.11.1:50090",776+ "npu": "tcp://10.0.11.1:50090",
779 "cpu": "tcp://10.0.11.1:50090",777 "cpu": "tcp://10.0.11.1:50090",
780 "disk": "tcp://10.0.11.1:50090"778 "disk": "tcp://10.0.11.1:50090"
781 },779 },
Mtests/controller/fault_tolerance/test_fault_manager.py+16-10
@@ -1925,17 +1925,23 @@ def test_pre_separate_l6_inactive_instances_triggers_scale_p2d(
1925 with patch(_FAULT_MGR_IM) as mock_fm_im_class:1925 with patch(_FAULT_MGR_IM) as mock_fm_im_class:
1926 mock_fm_im = MagicMock()1926 mock_fm_im = MagicMock()
1927 mock_fm_im_class.return_value = mock_fm_im1927 mock_fm_im_class.return_value = mock_fm_im
1928+ # Prevent async ScaleP2D.execute from racing: without a scale_p2d IM
1929+ # patch, execute aborts immediately, marks finished, and
1930+ # _process_instance_strategy clears strategy before assertions.
1931+ with patch.object(manager.executor, "submit") as mock_submit:
1932+ # First refresh fault level → L6 (included because node has instances)
1933+ manager._refresh_instance_fault_level(1)
1934+ # Then process strategy → should trigger ScaleP2D
1935+ manager._process_instance_strategy(1)
1928 1936 
1929- # First refresh fault level → L6 (included because node has instances)1937+ ins_meta = manager.instances[1]
1930- manager._refresh_instance_fault_level(1)1938+ assert ins_meta.fault_level == FaultLevel.L6, (
1931- # Then process strategy → should trigger ScaleP2D1939+ "PreSeparateNPU L6 with instances on node should set instance fault to L6"
1932- manager._process_instance_strategy(1)1940+ )
1933- 1941+ assert ins_meta.strategy is not None, (
1934- ins_meta = manager.instances[1]1942+ "PreSeparateNPU L6 with inactive instances should trigger ScaleP2D"
1935- assert ins_meta.fault_level == FaultLevel.L6, (1943+ )
1936- "PreSeparateNPU L6 with instances on node should set instance fault to L6"1944+ mock_submit.assert_called_once()
1937- )
1938- assert ins_meta.strategy is not None, "PreSeparateNPU L6 with inactive instances should trigger ScaleP2D"
1939 1945 
1940 1946 
1941# =============================================================================1947# =============================================================================
Mtests/coordinator/api_client/test_conductor_api_client.py+39-39
@@ -66,7 +66,7 @@ def _mock_config(**overrides) -> Mock:
66 66 
67 reg = KvConductorConfig(67 reg = KvConductorConfig(
68 store_backend=overrides.get("store_backend", "Mooncake"),68 store_backend=overrides.get("store_backend", "Mooncake"),
69- xpu_endpoint=overrides.get("xpu_endpoint", "tcp://*:5557"),69+ npu_endpoint=overrides.get("npu_endpoint", "tcp://*:5557"),
70 endpoint=overrides.get("endpoint", "tcp://*:5557"),70 endpoint=overrides.get("endpoint", "tcp://*:5557"),
71 replay_endpoint=overrides.get("replay_endpoint", ""),71 replay_endpoint=overrides.get("replay_endpoint", ""),
72 engine_type=overrides.get("engine_type", "vLLM"),72 engine_type=overrides.get("engine_type", "vLLM"),
@@ -122,7 +122,7 @@ class TestBuildRegisterPayload:
122 122 
123 def test_returns_empty_dict_when_no_endpoints_configured(self):123 def test_returns_empty_dict_when_no_endpoints_configured(self):
124 """No endpoint patterns configured → empty dict."""124 """No endpoint patterns configured → empty dict."""
125- cfg = _mock_config(xpu_endpoint="", endpoint="", replay_endpoint="")125+ cfg = _mock_config(npu_endpoint="", endpoint="", replay_endpoint="")
126 inst = _make_instance(inst_id=1, role=PDRole.ROLE_P)126 inst = _make_instance(inst_id=1, role=PDRole.ROLE_P)
127 ep = _make_endpoint(ep_id=0, ip="10.0.0.1")127 ep = _make_endpoint(ep_id=0, ip="10.0.0.1")
128 128 
@@ -131,9 +131,9 @@ class TestBuildRegisterPayload:
131 131 
132 assert payload == {}132 assert payload == {}
133 133 
134- def test_basic_payload_with_xpu_endpoint(self):134+ def test_basic_payload_with_npu_endpoint(self):
135- """Standard payload with medium_endpoints via xpu_endpoint (no fallback)."""135+ """Standard payload with medium_endpoints via npu_endpoint (no fallback)."""
136- cfg = _mock_config(xpu_endpoint="tcp://*:5557", endpoint="")136+ cfg = _mock_config(npu_endpoint="tcp://*:5557", endpoint="")
137 inst = _make_instance(inst_id=1, role=PDRole.ROLE_P, model_name="qwen")137 inst = _make_instance(inst_id=1, role=PDRole.ROLE_P, model_name="qwen")
138 ep = _make_endpoint(ep_id=0, ip="10.0.0.1")138 ep = _make_endpoint(ep_id=0, ip="10.0.0.1")
139 139 
@@ -142,7 +142,7 @@ class TestBuildRegisterPayload:
142 142 
143 assert payload["instance_id"] == "vllm-prefill-1"143 assert payload["instance_id"] == "vllm-prefill-1"
144 assert payload["dp_rank"] == 0144 assert payload["dp_rank"] == 0
145- assert payload["medium_endpoints"] == {"xpu": "tcp://10.0.0.1:5557"}145+ assert payload["medium_endpoints"] == {"npu": "tcp://10.0.0.1:5557"}
146 assert payload["type"] == "vLLM"146 assert payload["type"] == "vLLM"
147 assert payload["modelname"] == "qwen"147 assert payload["modelname"] == "qwen"
148 assert payload["block_size"] == 128148 assert payload["block_size"] == 128
@@ -150,7 +150,7 @@ class TestBuildRegisterPayload:
150 def test_payload_with_replay_endpoint(self):150 def test_payload_with_replay_endpoint(self):
151 """Payload includes replay_endpoint when configured."""151 """Payload includes replay_endpoint when configured."""
152 cfg = _mock_config(152 cfg = _mock_config(
153- xpu_endpoint="tcp://*:5557",153+ npu_endpoint="tcp://*:5557",
154 replay_endpoint="tcp://*:6667",154 replay_endpoint="tcp://*:6667",
155 )155 )
156 inst = _make_instance(inst_id=2, role=PDRole.ROLE_U, model_name="qwen")156 inst = _make_instance(inst_id=2, role=PDRole.ROLE_U, model_name="qwen")
@@ -165,7 +165,7 @@ class TestBuildRegisterPayload:
165 165 
166 def test_payload_dp_rank_uses_endpoint_id(self):166 def test_payload_dp_rank_uses_endpoint_id(self):
167 """dp_rank is taken from endpoint.id."""167 """dp_rank is taken from endpoint.id."""
168- cfg = _mock_config(xpu_endpoint="tcp://*:5557")168+ cfg = _mock_config(npu_endpoint="tcp://*:5557")
169 inst = _make_instance(inst_id=3, role=PDRole.ROLE_P)169 inst = _make_instance(inst_id=3, role=PDRole.ROLE_P)
170 ep = _make_endpoint(ep_id=5, ip="10.0.0.3")170 ep = _make_endpoint(ep_id=5, ip="10.0.0.3")
171 171 
@@ -173,25 +173,25 @@ class TestBuildRegisterPayload:
173 payload = ConductorApiClient._build_register_payload(inst, ep)173 payload = ConductorApiClient._build_register_payload(inst, ep)
174 174 
175 assert payload["dp_rank"] == 5175 assert payload["dp_rank"] == 5
176- assert payload["medium_endpoints"]["xpu"] == "tcp://10.0.0.3:5562"176+ assert payload["medium_endpoints"]["npu"] == "tcp://10.0.0.3:5562"
177 177 
178 def test_payload_with_fallback_endpoint(self):178 def test_payload_with_fallback_endpoint(self):
179- """Legacy 'endpoint' fallback pattern used when xpu_endpoint empty."""179+ """Legacy 'endpoint' fallback pattern used when npu_endpoint empty."""
180- cfg = _mock_config(xpu_endpoint="", endpoint="tcp://*:15557")180+ cfg = _mock_config(npu_endpoint="", endpoint="tcp://*:15557")
181 inst = _make_instance(inst_id=4, role=PDRole.ROLE_P)181 inst = _make_instance(inst_id=4, role=PDRole.ROLE_P)
182 ep = _make_endpoint(ep_id=0, ip="10.0.0.4")182 ep = _make_endpoint(ep_id=0, ip="10.0.0.4")
183 183 
184 with patch.object(ConductorApiClient, "coordinator_config", cfg):184 with patch.object(ConductorApiClient, "coordinator_config", cfg):
185 payload = ConductorApiClient._build_register_payload(inst, ep)185 payload = ConductorApiClient._build_register_payload(inst, ep)
186 186 
187- # Fallback endpoint fills xpu, cpu, disk187+ # Fallback endpoint fills gpu, cpu, disk
188 meps = payload["medium_endpoints"]188 meps = payload["medium_endpoints"]
189- assert "xpu" in meps189+ assert "npu" in meps
190 190 
191 def test_replay_endpoint_malformed_skipped(self):191 def test_replay_endpoint_malformed_skipped(self):
192 """replay_endpoint without '*:' → replay_endpoint absent in payload."""192 """replay_endpoint without '*:' → replay_endpoint absent in payload."""
193 cfg = _mock_config(193 cfg = _mock_config(
194- xpu_endpoint="tcp://*:5557",194+ npu_endpoint="tcp://*:5557",
195 replay_endpoint="tcp://127.0.0.1:6667",195 replay_endpoint="tcp://127.0.0.1:6667",
196 )196 )
197 inst = _make_instance(inst_id=5, role=PDRole.ROLE_P)197 inst = _make_instance(inst_id=5, role=PDRole.ROLE_P)
@@ -219,7 +219,7 @@ class TestNormalizeServiceKey:
219 "instance_id": "vllm-prefill-1",219 "instance_id": "vllm-prefill-1",
220 "endpoints": {220 "endpoints": {
221 "0": {221 "0": {
222- "medium_endpoints": {"xpu": "tcp://10.0.0.1:5557"},222+ "medium_endpoints": {"npu": "tcp://10.0.0.1:5557"},
223 "dp_rank": 0,223 "dp_rank": 0,
224 }224 }
225 },225 },
@@ -232,8 +232,8 @@ class TestNormalizeServiceKey:
232 worker = {232 worker = {
233 "instance_id": "vllm-union-2",233 "instance_id": "vllm-union-2",
234 "endpoints": {234 "endpoints": {
235- "0": {"medium_endpoints": {"xpu": "tcp://10.0.0.1:5557"}},235+ "0": {"medium_endpoints": {"npu": "tcp://10.0.0.1:5557"}},
236- "1": {"medium_endpoints": {"xpu": "tcp://10.0.0.1:5558"}},236+ "1": {"medium_endpoints": {"npu": "tcp://10.0.0.1:5558"}},
237 },237 },
238 }238 }
239 keys = ConductorApiClient._normalize_service_key(worker)239 keys = ConductorApiClient._normalize_service_key(worker)
@@ -250,8 +250,8 @@ class TestNormalizeServiceKey:
250 worker = {250 worker = {
251 "instance_id": "vllm-prefill-1",251 "instance_id": "vllm-prefill-1",
252 "endpoints": {252 "endpoints": {
253- "abc": {"medium_endpoints": {"xpu": "tcp://x:1"}},253+ "abc": {"medium_endpoints": {"npu": "tcp://x:1"}},
254- "0": {"medium_endpoints": {"xpu": "tcp://x:2"}},254+ "0": {"medium_endpoints": {"npu": "tcp://x:2"}},
255 },255 },
256 }256 }
257 keys = ConductorApiClient._normalize_service_key(worker)257 keys = ConductorApiClient._normalize_service_key(worker)
@@ -419,7 +419,7 @@ class TestReRegisterKvInstances:
419 def test_skip_when_no_endpoints_configured(self):419 def test_skip_when_no_endpoints_configured(self):
420 """No endpoint patterns → _build_register_payload returns {} → skip."""420 """No endpoint patterns → _build_register_payload returns {} → skip."""
421 inst = _make_instance(inst_id=1, role=PDRole.ROLE_P)421 inst = _make_instance(inst_id=1, role=PDRole.ROLE_P)
422- cfg = _mock_config(xpu_endpoint="", endpoint="")422+ cfg = _mock_config(npu_endpoint="", endpoint="")
423 423 
424 with (424 with (
425 patch.object(ConductorApiClient, "coordinator_config", cfg),425 patch.object(ConductorApiClient, "coordinator_config", cfg),
@@ -436,7 +436,7 @@ class TestReRegisterKvInstances:
436 endpoints = {"pod-0": {0: ep}}436 endpoints = {"pod-0": {0: ep}}
437 inst = Instance(id=1, role=PDRole.ROLE_P, model_name="qwen", job_name="test-job", endpoints=endpoints)437 inst = Instance(id=1, role=PDRole.ROLE_P, model_name="qwen", job_name="test-job", endpoints=endpoints)
438 438 
439- cfg = _mock_config(xpu_endpoint="tcp://*:5557")439+ cfg = _mock_config(npu_endpoint="tcp://*:5557")
440 440 
441 # Conductor has a DIFFERENT instance registered441 # Conductor has a DIFFERENT instance registered
442 registered = [{"instance_id": "vllm-prefill-99", "endpoints": {"0": {}}}]442 registered = [{"instance_id": "vllm-prefill-99", "endpoints": {"0": {}}}]
@@ -456,11 +456,11 @@ class TestReRegisterKvInstances:
456 endpoints = {"pod-0": {0: ep}}456 endpoints = {"pod-0": {0: ep}}
457 inst = Instance(id=1, role=PDRole.ROLE_P, model_name="qwen", job_name="test-job", endpoints=endpoints)457 inst = Instance(id=1, role=PDRole.ROLE_P, model_name="qwen", job_name="test-job", endpoints=endpoints)
458 458 
459- cfg = _mock_config(xpu_endpoint="tcp://*:5557")459+ cfg = _mock_config(npu_endpoint="tcp://*:5557")
460 460 
461 # Same (instance_id, dp_rank) already registered461 # Same (instance_id, dp_rank) already registered
462 registered = [462 registered = [
463- {"instance_id": "vllm-prefill-1", "endpoints": {"0": {"medium_endpoints": {"xpu": "tcp://10.0.0.1:5557"}}}}463+ {"instance_id": "vllm-prefill-1", "endpoints": {"0": {"medium_endpoints": {"npu": "tcp://10.0.0.1:5557"}}}}
464 ]464 ]
465 465 
466 with (466 with (
@@ -479,11 +479,11 @@ class TestReRegisterKvInstances:
479 endpoints = {"pod-0": {0: ep0, 1: ep1}}479 endpoints = {"pod-0": {0: ep0, 1: ep1}}
480 inst = Instance(id=1, role=PDRole.ROLE_P, model_name="qwen", job_name="test-job", endpoints=endpoints)480 inst = Instance(id=1, role=PDRole.ROLE_P, model_name="qwen", job_name="test-job", endpoints=endpoints)
481 481 
482- cfg = _mock_config(xpu_endpoint="tcp://*:5557")482+ cfg = _mock_config(npu_endpoint="tcp://*:5557")
483 483 
484 # ep0 (dp_rank=0) already registered; ep1 (dp_rank=1) missing484 # ep0 (dp_rank=0) already registered; ep1 (dp_rank=1) missing
485 registered = [485 registered = [
486- {"instance_id": "vllm-prefill-1", "endpoints": {"0": {"medium_endpoints": {"xpu": "tcp://10.0.0.1:5557"}}}}486+ {"instance_id": "vllm-prefill-1", "endpoints": {"0": {"medium_endpoints": {"npu": "tcp://10.0.0.1:5557"}}}}
487 ]487 ]
488 488 
489 with (489 with (
@@ -504,7 +504,7 @@ class TestReRegisterKvInstances:
504 endpoints = {"pod-0": {0: ep}}504 endpoints = {"pod-0": {0: ep}}
505 inst = Instance(id=1, role=PDRole.ROLE_P, model_name="qwen", job_name="test-job", endpoints=endpoints)505 inst = Instance(id=1, role=PDRole.ROLE_P, model_name="qwen", job_name="test-job", endpoints=endpoints)
506 506 
507- cfg = _mock_config(xpu_endpoint="tcp://*:5557")507+ cfg = _mock_config(npu_endpoint="tcp://*:5557")
508 508 
509 # Mooncake Master format: InstanceID + DPRank509 # Mooncake Master format: InstanceID + DPRank
510 registered = [{"InstanceID": "vllm-prefill-1", "DPRank": 0, "Endpoint": "tcp://10.0.0.1:5557"}]510 registered = [{"InstanceID": "vllm-prefill-1", "DPRank": 0, "Endpoint": "tcp://10.0.0.1:5557"}]
@@ -524,7 +524,7 @@ class TestReRegisterKvInstances:
524 endpoints = {"pod-0": {0: ep}}524 endpoints = {"pod-0": {0: ep}}
525 inst = Instance(id=1, role=PDRole.ROLE_P, model_name="qwen", job_name="test-job", endpoints=endpoints)525 inst = Instance(id=1, role=PDRole.ROLE_P, model_name="qwen", job_name="test-job", endpoints=endpoints)
526 526 
527- cfg = _mock_config(xpu_endpoint="tcp://*:5557")527+ cfg = _mock_config(npu_endpoint="tcp://*:5557")
528 528 
529 # Different instance registered529 # Different instance registered
530 registered = [{"InstanceID": "vllm-prefill-99", "DPRank": 0, "Endpoint": "tcp://10.0.0.99:5557"}]530 registered = [{"InstanceID": "vllm-prefill-99", "DPRank": 0, "Endpoint": "tcp://10.0.0.99:5557"}]
@@ -593,7 +593,7 @@ def test_return_value_on_failure(mock_http):
593 593 
594 594 
595def _setup_reg_config(595def _setup_reg_config(
596- store_backend, pool_endpoint="", xpu_endpoint="", cpu_endpoint="", disk_endpoint="", replay_endpoint=""596+ store_backend, pool_endpoint="", npu_endpoint="", cpu_endpoint="", disk_endpoint="", replay_endpoint=""
597):597):
598 """Patch ConductorApiClient's config for registration testing."""598 """Patch ConductorApiClient's config for registration testing."""
599 from motor.config.coordinator import KvConductorConfig, SchedulerConfig599 from motor.config.coordinator import KvConductorConfig, SchedulerConfig
@@ -601,7 +601,7 @@ def _setup_reg_config(
601 reg = KvConductorConfig(601 reg = KvConductorConfig(
602 store_backend=store_backend,602 store_backend=store_backend,
603 pool_endpoint=pool_endpoint,603 pool_endpoint=pool_endpoint,
604- xpu_endpoint=xpu_endpoint,604+ npu_endpoint=npu_endpoint,
605 cpu_endpoint=cpu_endpoint,605 cpu_endpoint=cpu_endpoint,
606 disk_endpoint=disk_endpoint,606 disk_endpoint=disk_endpoint,
607 replay_endpoint=replay_endpoint,607 replay_endpoint=replay_endpoint,
@@ -634,7 +634,7 @@ def test_yuanrong_registration_dispatches_per_dp(mock_http):
634 ConductorApiClient._pool_registered = False634 ConductorApiClient._pool_registered = False
635 635 
636 with _setup_reg_config(636 with _setup_reg_config(
637- "YuanRong", xpu_endpoint="tcp://*:15557", cpu_endpoint="tcp://*:15558", disk_endpoint="tcp://*:15558"637+ "YuanRong", npu_endpoint="tcp://*:15557", cpu_endpoint="tcp://*:15558", disk_endpoint="tcp://*:15558"
638 ):638 ):
639 ConductorApiClient.register_kv_instance([instance])639 ConductorApiClient.register_kv_instance([instance])
640 640 
@@ -643,7 +643,7 @@ def test_yuanrong_registration_dispatches_per_dp(mock_http):
643 payload = calls[0][0][1]643 payload = calls[0][0][1]
644 assert "medium_endpoints" in payload644 assert "medium_endpoints" in payload
645 assert payload["store_backend"] == "YuanRong"645 assert payload["store_backend"] == "YuanRong"
646- assert "xpu" in str(payload["medium_endpoints"])646+ assert "npu" in str(payload["medium_endpoints"])
647 assert "cpu" in str(payload["medium_endpoints"])647 assert "cpu" in str(payload["medium_endpoints"])
648 assert "disk" in str(payload["medium_endpoints"])648 assert "disk" in str(payload["medium_endpoints"])
649 649 
@@ -658,7 +658,7 @@ def test_mooncake_registration_includes_pool_plus_hbm(mock_http):
658 instance = _make_mock_instance(1)658 instance = _make_mock_instance(1)
659 ConductorApiClient._pool_registered = False659 ConductorApiClient._pool_registered = False
660 660 
661- with _setup_reg_config("Mooncake", pool_endpoint="tcp://kvp-master:5557", xpu_endpoint="tcp://*:50090"):661+ with _setup_reg_config("Mooncake", pool_endpoint="tcp://kvp-master:5557", npu_endpoint="tcp://*:50090"):
662 ConductorApiClient.register_kv_instance([instance])662 ConductorApiClient.register_kv_instance([instance])
663 663 
664 calls = mock_client.post.call_args_list664 calls = mock_client.post.call_args_list
@@ -673,7 +673,7 @@ def test_mooncake_registration_includes_pool_plus_hbm(mock_http):
673 # Second call: HBM DP673 # Second call: HBM DP
674 hbm_payload = calls[1][0][1]674 hbm_payload = calls[1][0][1]
675 assert "medium_endpoints" in hbm_payload675 assert "medium_endpoints" in hbm_payload
676- assert "xpu" in str(hbm_payload["medium_endpoints"])676+ assert "npu" in str(hbm_payload["medium_endpoints"])
677 677 
678 678 
679@patch("motor.coordinator.api_client.conductor_api_client.SafeHTTPSClient")679@patch("motor.coordinator.api_client.conductor_api_client.SafeHTTPSClient")
@@ -686,7 +686,7 @@ def test_mooncake_pool_only_registered_once(mock_http):
686 instance = _make_mock_instance(1)686 instance = _make_mock_instance(1)
687 ConductorApiClient._pool_registered = False687 ConductorApiClient._pool_registered = False
688 688 
689- with _setup_reg_config("Mooncake", pool_endpoint="tcp://kvp-master:5557", xpu_endpoint="tcp://*:50090"):689+ with _setup_reg_config("Mooncake", pool_endpoint="tcp://kvp-master:5557", npu_endpoint="tcp://*:50090"):
690 ConductorApiClient.register_kv_instance([instance])690 ConductorApiClient.register_kv_instance([instance])
691 ConductorApiClient.register_kv_instance([instance])691 ConductorApiClient.register_kv_instance([instance])
692 692 
@@ -707,7 +707,7 @@ def test_memcache_registration_same_as_mooncake_different_store_backend(mock_htt
707 instance = _make_mock_instance(1)707 instance = _make_mock_instance(1)
708 ConductorApiClient._pool_registered = False708 ConductorApiClient._pool_registered = False
709 709 
710- with _setup_reg_config("Memcache", pool_endpoint="tcp://kvp-master:5557", xpu_endpoint="tcp://*:50090"):710+ with _setup_reg_config("Memcache", pool_endpoint="tcp://kvp-master:5557", npu_endpoint="tcp://*:50090"):
711 ConductorApiClient.register_kv_instance([instance])711 ConductorApiClient.register_kv_instance([instance])
712 712 
713 calls = mock_client.post.call_args_list713 calls = mock_client.post.call_args_list
@@ -728,7 +728,7 @@ def test_replay_endpoint_included_in_registration(mock_http):
728 728 
729 with _setup_reg_config(729 with _setup_reg_config(
730 "YuanRong",730 "YuanRong",
731- xpu_endpoint="tcp://*:15557",731+ npu_endpoint="tcp://*:15557",
732 cpu_endpoint="tcp://*:15558",732 cpu_endpoint="tcp://*:15558",
733 disk_endpoint="tcp://*:15558",733 disk_endpoint="tcp://*:15558",
734 replay_endpoint="tcp://*:6667",734 replay_endpoint="tcp://*:6667",
@@ -752,13 +752,13 @@ def test_endpoint_url_resolves_ip_and_dp_rank(mock_http):
752 instance.endpoints["pod-0"][0].id = 2 # dp_rank=2752 instance.endpoints["pod-0"][0].id = 2 # dp_rank=2
753 753 
754 with _setup_reg_config(754 with _setup_reg_config(
755- "YuanRong", xpu_endpoint="tcp://*:15557", cpu_endpoint="tcp://*:15558", disk_endpoint="tcp://*:15558"755+ "YuanRong", npu_endpoint="tcp://*:15557", cpu_endpoint="tcp://*:15558", disk_endpoint="tcp://*:15558"
756 ):756 ):
757 ConductorApiClient.register_kv_instance([instance])757 ConductorApiClient.register_kv_instance([instance])
758 758 
759 payload = mock_client.post.call_args_list[0][0][1]759 payload = mock_client.post.call_args_list[0][0][1]
760 meps = payload["medium_endpoints"]760 meps = payload["medium_endpoints"]
761- assert meps["xpu"] == "tcp://10.0.0.1:15559" # 15557 + 2761+ assert meps["npu"] == "tcp://10.0.0.1:15559" # 15557 + 2
762 assert meps["cpu"] == "tcp://10.0.0.1:15560" # 15558 + 2762 assert meps["cpu"] == "tcp://10.0.0.1:15560" # 15558 + 2
763 assert payload["dp_rank"] == 2763 assert payload["dp_rank"] == 2
764 764 
@@ -772,7 +772,7 @@ def test_unknown_backend_falls_back_to_per_dp(mock_http):
772 772 
773 instance = _make_mock_instance(1)773 instance = _make_mock_instance(1)
774 774 
775- with _setup_reg_config("SomeUnknownBackend", xpu_endpoint="tcp://*:15557"):775+ with _setup_reg_config("SomeUnknownBackend", npu_endpoint="tcp://*:15557"):
776 ConductorApiClient.register_kv_instance([instance])776 ConductorApiClient.register_kv_instance([instance])
777 777 
778 assert len(mock_client.post.call_args_list) >= 1778 assert len(mock_client.post.call_args_list) >= 1
Mtests/coordinator/scheduler/test_kv_cache_affinity.py+4-4
@@ -268,7 +268,7 @@ class TestKvCacheAffinityPolicy(unittest.TestCase):
268 @patch('motor.coordinator.scheduler.policy.kv_cache_affinity.ConductorApiClient.query_conductor')268 @patch('motor.coordinator.scheduler.policy.kv_cache_affinity.ConductorApiClient.query_conductor')
269 @patch('motor.coordinator.scheduler.policy.kv_cache_affinity.TokenizerManager')269 @patch('motor.coordinator.scheduler.policy.kv_cache_affinity.TokenizerManager')
270 def test_select_endpoint_dpscoring_dict_format(self, mock_tokenizer_manager, mock_query_conductor):270 def test_select_endpoint_dpscoring_dict_format(self, mock_tokenizer_manager, mock_query_conductor):
271- """New DpScoring format: DP values are dicts with matched_tokens, not plain ints."""271+ """New DpBlocks format: DP values are dicts with matched_tokens, not plain ints."""
272 ep_a = _make_endpoint(0, active_tokens=50.0)272 ep_a = _make_endpoint(0, active_tokens=50.0)
273 ep_b = _make_endpoint(1, active_tokens=50.0)273 ep_b = _make_endpoint(1, active_tokens=50.0)
274 mock_instance = Mock()274 mock_instance = Mock()
@@ -284,13 +284,13 @@ class TestKvCacheAffinityPolicy(unittest.TestCase):
284 mock_tokenizer.encode.return_value = list(range(1000))284 mock_tokenizer.encode.return_value = list(range(1000))
285 mock_tokenizer_manager.return_value = mock_tokenizer285 mock_tokenizer_manager.return_value = mock_tokenizer
286 286 
287- # New DpScoring format: DP values are dicts with matched_tokens287+ # New DpBlocks format: DP values are dicts with matched_tokens
288 mock_query_conductor.return_value = {288 mock_query_conductor.return_value = {
289 TENANT_ID: {289 TENANT_ID: {
290 "vllm-prefill-inst": {290 "vllm-prefill-inst": {
291 "DP": {291 "DP": {
292- "0": {"XPU": 2400, "CPU": 0, "DISK": 0, "total": 2400, "matched_tokens": 800},292+ "0": {"npu_blocks": 6, "cpu_blocks": 0, "disk_blocks": 0, "matched_tokens": 800},
293- "1": {"XPU": 0, "CPU": 200, "DISK": 0, "total": 200, "matched_tokens": 100},293+ "1": {"npu_blocks": 0, "cpu_blocks": 1, "disk_blocks": 0, "matched_tokens": 100},
294 }294 }
295 }295 }
296 }296 }
Mtests/coordinator/scheduler/test_kva_role_u_support.py+4-2
@@ -68,6 +68,7 @@ def test_register_post_uses_union_conductor_id_for_role_u() -> None:
68 mock_config.scheduler_config.kv_conductor_config.conductor_service = "kv-conductor"68 mock_config.scheduler_config.kv_conductor_config.conductor_service = "kv-conductor"
69 mock_config.scheduler_config.kv_conductor_config.http_server_port = 1333369 mock_config.scheduler_config.kv_conductor_config.http_server_port = 13333
70 mock_config.scheduler_config.kv_conductor_config.endpoint = "tcp://*:5557"70 mock_config.scheduler_config.kv_conductor_config.endpoint = "tcp://*:5557"
71+ mock_config.scheduler_config.kv_conductor_config.npu_endpoint = ""
71 mock_config.scheduler_config.kv_conductor_config.xpu_endpoint = ""72 mock_config.scheduler_config.kv_conductor_config.xpu_endpoint = ""
72 mock_config.scheduler_config.kv_conductor_config.cpu_endpoint = ""73 mock_config.scheduler_config.kv_conductor_config.cpu_endpoint = ""
73 mock_config.scheduler_config.kv_conductor_config.disk_endpoint = ""74 mock_config.scheduler_config.kv_conductor_config.disk_endpoint = ""
@@ -101,6 +102,7 @@ def test_register_post_formats_ipv6_endpoint_and_conductor_address() -> None:
101 mock_config.scheduler_config.kv_conductor_config.model_path = ""102 mock_config.scheduler_config.kv_conductor_config.model_path = ""
102 mock_config.scheduler_config.kv_conductor_config.conductor_service = "2001:db8::10"103 mock_config.scheduler_config.kv_conductor_config.conductor_service = "2001:db8::10"
103 mock_config.scheduler_config.kv_conductor_config.http_server_port = 13333104 mock_config.scheduler_config.kv_conductor_config.http_server_port = 13333
105+ mock_config.scheduler_config.kv_conductor_config.npu_endpoint = ""
104 mock_config.scheduler_config.kv_conductor_config.xpu_endpoint = ""106 mock_config.scheduler_config.kv_conductor_config.xpu_endpoint = ""
105 mock_config.scheduler_config.kv_conductor_config.cpu_endpoint = ""107 mock_config.scheduler_config.kv_conductor_config.cpu_endpoint = ""
106 mock_config.scheduler_config.kv_conductor_config.disk_endpoint = ""108 mock_config.scheduler_config.kv_conductor_config.disk_endpoint = ""
@@ -118,9 +120,9 @@ def test_register_post_formats_ipv6_endpoint_and_conductor_address() -> None:
118 mock_http_client.assert_called_once()120 mock_http_client.assert_called_once()
119 assert mock_http_client.call_args.kwargs["address"] == "[2001:db8::10]:13333"121 assert mock_http_client.call_args.kwargs["address"] == "[2001:db8::10]:13333"
120 register_payload = mock_http_client.return_value.__enter__.return_value.post.call_args[0][1]122 register_payload = mock_http_client.return_value.__enter__.return_value.post.call_args[0][1]
121- # Endpoints are now wrapped in medium_endpoints dict (xpu/cpu/disk).123+ # Endpoints are now wrapped in medium_endpoints dict (npu/cpu/disk).
122 # When per-medium fields are empty, all fall back to the legacy "endpoint".124 # When per-medium fields are empty, all fall back to the legacy "endpoint".
123- assert register_payload["medium_endpoints"]["xpu"] == "tcp://[2001:db8::1]:5559"125+ assert register_payload["medium_endpoints"]["npu"] == "tcp://[2001:db8::1]:5559"
124 assert register_payload["replay_endpoint"] == "tcp://[2001:db8::1]:6669"126 assert register_payload["replay_endpoint"] == "tcp://[2001:db8::1]:6669"
125 127 
126 128