dynamo_mocker/common/
kv_cache_trace.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Shared KV cache trace logging for both vLLM and SGLang backends.
5//!
6//! Enabled by setting `DYN_MOCKER_KV_CACHE_TRACE=1` or `true`.
7
8use std::sync::LazyLock;
9use std::time::{SystemTime, UNIX_EPOCH};
10
11const DYN_MOCKER_KV_CACHE_TRACE: &str = "DYN_MOCKER_KV_CACHE_TRACE";
12
13/// Check the env var to enable KV cache allocation/eviction trace logs.
14pub static KV_CACHE_TRACE_ENABLED: LazyLock<bool> =
15    LazyLock::new(|| dynamo_truthy::env_is_truthy(DYN_MOCKER_KV_CACHE_TRACE));
16
17fn timestamp_ms() -> u64 {
18    SystemTime::now()
19        .duration_since(UNIX_EPOCH)
20        .unwrap_or_default()
21        .as_millis() as u64
22}
23
24/// Log a vLLM KV cache trace event.
25pub fn log_vllm_trace(
26    event: &str,
27    dp_rank: u32,
28    block_size: usize,
29    active_blocks: usize,
30    inactive_blocks: usize,
31    total_blocks: usize,
32) {
33    if !*KV_CACHE_TRACE_ENABLED {
34        return;
35    }
36    let free_blocks = total_blocks
37        .saturating_sub(active_blocks)
38        .saturating_sub(inactive_blocks);
39    let utilization = if total_blocks > 0 {
40        (active_blocks + inactive_blocks) as f64 / total_blocks as f64
41    } else {
42        0.0
43    };
44    tracing::info!(
45        engine_type = "vllm",
46        event,
47        timestamp_ms = timestamp_ms(),
48        dp_rank,
49        block_size,
50        free_blocks,
51        active_blocks,
52        inactive_blocks,
53        total_blocks,
54        utilization,
55        "KV cache trace"
56    );
57}
58
59/// SGLang cache state snapshot for trace logging.
60pub struct SglangCacheState<'a> {
61    pub event: &'a str,
62    pub dp_rank: u32,
63    pub num_tokens: usize,
64    pub page_size: usize,
65    pub available_tokens: usize,
66    pub evictable_tokens: usize,
67    pub protected_tokens: usize,
68    pub total_tokens: usize,
69}
70
71/// Log an SGLang KV cache trace event.
72pub fn log_sglang_trace(state: &SglangCacheState) {
73    if !*KV_CACHE_TRACE_ENABLED {
74        return;
75    }
76    let utilization = if state.total_tokens > 0 {
77        (state.total_tokens - state.available_tokens) as f64 / state.total_tokens as f64
78    } else {
79        0.0
80    };
81    tracing::info!(
82        engine_type = "sglang",
83        event = state.event,
84        timestamp_ms = timestamp_ms(),
85        dp_rank = state.dp_rank,
86        num_tokens = state.num_tokens,
87        page_size = state.page_size,
88        available_tokens = state.available_tokens,
89        evictable_tokens = state.evictable_tokens,
90        protected_tokens = state.protected_tokens,
91        total_tokens = state.total_tokens,
92        utilization,
93        "KV cache trace"
94    );
95}
dynamo_mocker/common/kv_cache_trace.rs

dynamo_mocker/common/
kv_cache_trace.rs
