fix(pylon): exclude cached prompt tokens from fallback TPS

src/libraries/rust/stargate/crates/mock-dynamo/src/openai.rs

@@ -159,10 +159,16 @@ struct AssistantMessage<'a> {
 #[derive(Serialize)]
 struct ChatUsage {
     prompt_tokens: usize,
+    prompt_tokens_details: ChatPromptTokensDetails,
     completion_tokens: usize,
     total_tokens: usize,
 }
 
+#[derive(Serialize)]
+struct ChatPromptTokensDetails {
+    cached_tokens: u64,
+}
+
 #[derive(Serialize)]
 #[serde(untagged)]
 pub(crate) enum EmbeddingValue {
@@ -351,6 +357,9 @@ pub(crate) async fn chat_completions(
         }],
         usage: ChatUsage {
             prompt_tokens: input_tokens,
+            prompt_tokens_details: ChatPromptTokensDetails {
+                cached_tokens: kv_cache_access.reused_input_tokens,
+            },
             completion_tokens: output_tokens,
             total_tokens: input_tokens.saturating_add(output_tokens),
         },
@@ -608,6 +617,7 @@ pub(crate) enum ChatStreamChunk<'a> {
     Stop,
     Usage {
         input_tokens: usize,
+        cached_input_tokens: u64,
         output_tokens: usize,
     },
 }
@@ -625,10 +635,14 @@ pub(crate) fn chat_chunk_json(
         ChatStreamChunk::Stop => Some((None, None, Some("stop"))),
         ChatStreamChunk::Usage {
             input_tokens,
+            cached_input_tokens,
             output_tokens,
         } => {
             usage = Some(Some(ChatUsage {
                 prompt_tokens: input_tokens,
+                prompt_tokens_details: ChatPromptTokensDetails {
+                    cached_tokens: cached_input_tokens,
+                },
                 completion_tokens: output_tokens,
                 total_tokens: input_tokens.saturating_add(output_tokens),
             }));
@@ -666,6 +680,7 @@ fn stream_response(config: StreamResponseConfig) -> Response {
         mut pacing,
         kind,
     } = config;
+    let cached_input_tokens = kv_cache_access.reused_input_tokens;
     let stream = async_stream::stream! {
         let mut output_text = String::new();
         if let StreamKind::Responses { created_at } = kind {
@@ -747,6 +762,9 @@ fn stream_response(config: StreamResponseConfig) -> Response {
                         }],
                         "usage": {
                             "input_tokens": input_tokens,
+                            "input_tokens_details": {
+                                "cached_tokens": cached_input_tokens,
+                            },
                             "output_tokens": output_tokens,
                             "total_tokens": input_tokens.saturating_add(output_tokens),
                         },
@@ -760,7 +778,7 @@ fn stream_response(config: StreamResponseConfig) -> Response {
 
         if let StreamKind::Chat { include_usage, .. } = kind {
             if include_usage {
-                yield Ok(chat_sse_event(&id, &model, ChatStreamChunk::Usage { input_tokens, output_tokens }, true));
+                yield Ok(chat_sse_event(&id, &model, ChatStreamChunk::Usage { input_tokens, cached_input_tokens, output_tokens }, true));
             }
             yield Ok(Event::default().data("[DONE]"));
         }

src/libraries/rust/stargate/crates/mock-dynamo/src/tests.rs

@@ -414,7 +414,7 @@ async fn test_control_http_api_updates_one_model_and_reports_request_counters()
     assert_eq!(
         body,
         serde_json::from_str::<serde_json::Value>(
-            r#"{"id":"chatcmpl-mock-id","object":"chat.completion","model":"model-a","choices":[{"index":0,"message":{"role":"assistant","content":"Hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}"#,
+            r#"{"id":"chatcmpl-mock-id","object":"chat.completion","model":"model-a","choices":[{"index":0,"message":{"role":"assistant","content":"Hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"prompt_tokens_details":{"cached_tokens":0},"completion_tokens":1,"total_tokens":2}}"#,
         )
         .expect("expected chat response fixture should be valid JSON")
     );
@@ -597,6 +597,7 @@ async fn streaming_chat_usage_reports_actual_output_only_when_requested() {
             distribution: OutputTokenDistribution::Uniform,
         },
         context_length_tokens: 5,
+        kv_cache: Arc::new(Mutex::new(KvCacheState::new(10))),
         ..test_state()
     };
     let app = Router::new()
@@ -614,7 +615,7 @@ async fn streaming_chat_usage_reports_actual_output_only_when_requested() {
             address,
             "POST",
             "/v1/chat/completions",
-            "connection: close\r\nx-input-tokens: 2\r\nx-output-tokens: 100",
+            "connection: close\r\nx-input-tokens: 2\r\nx-output-tokens: 100\r\nx-cache-affinity-key: chat-usage",
             &body.to_string(),
         )
         .await;
@@ -640,7 +641,12 @@ async fn streaming_chat_usage_reports_actual_output_only_when_requested() {
             assert_eq!(usage["choices"], serde_json::json!([]));
             assert_eq!(
                 usage["usage"],
-                serde_json::json!({"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5})
+                serde_json::json!({
+                    "prompt_tokens": 2,
+                    "prompt_tokens_details": {"cached_tokens": 2},
+                    "completion_tokens": 3,
+                    "total_tokens": 5
+                })
             );
             assert!(
                 output_events
@@ -1137,13 +1143,14 @@ async fn responses_endpoint_streams_response_events_without_private_stats_header
     let (addr, server) = spawn_test_app(app).await;
 
     let cases = [
-        (None, 1),
-        (Some(serde_json::json!("hello")), 5),
-        (Some(serde_json::json!(["hi", {"x": 1}])), 9),
-        (Some(serde_json::json!({"x": 1})), 7),
+        (None, 1, 1),
+        (Some(serde_json::json!("hello")), 5, 4),
+        (Some(serde_json::json!(["hi", {"x": 1}])), 9, 4),
+        (Some(serde_json::json!({"x": 1})), 7, 0),
     ];
     let mut response = String::new();
-    for (index, (input, expected_tokens)) in cases.into_iter().enumerate() {
+    for (index, (input, expected_tokens, expected_uncached_tokens)) in cases.into_iter().enumerate()
+    {
         let mut request_body = serde_json::json!({
             "model": "request-model",
             "max_output_tokens": 2,
@@ -1156,19 +1163,19 @@ async fn responses_endpoint_streams_response_events_without_private_stats_header
             addr,
             "POST",
             "/v1/responses",
-            &format!("connection: close\r\nx-request-id: req-responses-{index}\r\nx-input-tokens: {expected_tokens}\r\nx-cache-affinity-key: cache-{index}"),
+            &format!("connection: close\r\nx-request-id: req-responses-{index}\r\nx-input-tokens: {expected_tokens}\r\nx-cache-affinity-key: cache-shared"),
             &request_body.to_string(),
         )
         .await;
         assert!(response.contains(&format!(
-            "x-kv-cache-uncached-input-tokens: {expected_tokens}"
+            "x-kv-cache-uncached-input-tokens: {expected_uncached_tokens}"
         )));
         assert!(response.contains(&format!(r#""input_tokens":{expected_tokens}"#)));
     }
     assert!(response.starts_with("HTTP/1.1 200 OK"));
-    assert!(response.contains("x-kv-cache-hit: false"));
-    assert!(response.contains("x-kv-cache-reused-input-tokens: 0"));
-    assert!(response.contains("x-kv-cache-uncached-input-tokens: 7"));
+    assert!(response.contains("x-kv-cache-hit: true"));
+    assert!(response.contains("x-kv-cache-reused-input-tokens: 7"));
+    assert!(response.contains("x-kv-cache-uncached-input-tokens: 0"));
     assert!(response.contains("content-type: text/event-stream"));
 
     let (_, body_text) = response
@@ -1182,6 +1189,7 @@ async fn responses_endpoint_streams_response_events_without_private_stats_header
     assert!(body_text.contains(r#""status":"completed""#));
     assert!(body_text.contains(r#""model":"request-model""#));
     assert!(body_text.contains(r#""input_tokens":7"#));
+    assert!(body_text.contains(r#""input_tokens_details":{"cached_tokens":7}"#));
     assert!(body_text.contains(r#""output_tokens":2"#));
     assert!(body_text.contains(r#""total_tokens":9"#));
 

src/libraries/rust/stargate/crates/proto/proto/stargate.proto

@@ -123,8 +123,9 @@ message ModelStats {
   // Sticky source labels for stat observations seen by this backend since the
   // model metrics state was initialized.
   repeated string stats_sources = 17;
-  // Greatest last_mean_input_tps published for the current backend model
-  // generation. Absent until the first mean is configured or observed.
+  // Greatest max-eligible last_mean_input_tps published for the current backend
+  // model generation. Means that may count cached prompt tokens as prefill do
+  // not raise it. Absent until the first eligible mean is configured or observed.
   optional double max_input_tps = 18;
 }