[Feature] Add process-local in-memory KV indexer and Router integration (#33370)

Co-authored-by: Wu, Yutong <yutong.wu@amd.com>
Co-authored-by: TianDi101 <ditian12@amd.com>
Co-authored-by: Zhangheng <hzh0425@apache.org>
This commit is contained in:
wuyl1
2026-08-20 10:45:35 +08:00
committed by GitHub
co-authored by Wu, Yutong TianDi101 Zhangheng
parent 238ba40c27
commit 360d10d6bc
42 changed files with 6603 additions and 174 deletions
@@ -0,0 +1,57 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use sgl_kv_indexer::pb::{
ApplyExternalKvBatchRequest, ExternalKvAction, ExternalKvActionType, TierType,
};
pub fn hbm() -> i32 {
TierType::TierHbm as i32
}
pub fn dram() -> i32 {
TierType::TierDram as i32
}
pub fn action(kind: ExternalKvActionType, tier: i32, hashes: &[i64]) -> ExternalKvAction {
ExternalKvAction {
r#type: kind as i32,
tier,
hashes: hashes.to_vec(),
component_masks: Vec::new(),
block_sizes: Vec::new(),
}
}
/// A component-aware REPORT action: each hash carries its component bitmask and
/// token count, index-aligned with `hashes`.
#[allow(dead_code)] // used by memory_integration, not grpc_contract
pub fn component_report(
tier: i32,
hashes: &[i64],
masks: &[u32],
block_sizes: &[u32],
) -> ExternalKvAction {
ExternalKvAction {
r#type: ExternalKvActionType::ActionReport as i32,
tier,
hashes: hashes.to_vec(),
component_masks: masks.to_vec(),
block_sizes: block_sizes.to_vec(),
}
}
pub fn apply_request(
worker: &str,
address: &str,
seq: u64,
actions: Vec<ExternalKvAction>,
) -> ApplyExternalKvBatchRequest {
ApplyExternalKvBatchRequest {
worker_id: worker.to_string(),
seq,
actions,
worker_address: address.to_string(),
cache_spec: None,
}
}