Add kv cache hit source breakdown (device/host/storage) to step profiling

ThunderAgent/app.py

@@ -84,6 +84,7 @@ async def on_usage(
         prompt_tokens: int | None,
         completion_tokens: int | None,
         cached_tokens: int | None,
+        cached_tokens_details: Dict[str, int | str],
     ) -> None:
         router.update_program_after_request(
             program_id,
@@ -97,6 +98,7 @@ async def on_usage(
                 prompt_tokens=prompt_tokens,
                 completion_tokens=completion_tokens,
                 cached_tokens=cached_tokens,
+                cached_tokens_details=cached_tokens_details,
             )
 
     # Callback for streaming token progress updates (every 20 tokens)

ThunderAgent/profile/state.py

@@ -3,7 +3,7 @@
 import time
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 
 def get_profile_csv_dir() -> Path:
@@ -25,6 +25,10 @@ class StepMetrics:
     completion_tokens: Optional[int] = None
     cached_tokens: Optional[int] = None
     kv_hit_rate: Optional[float] = None  # KV cache hit rate (0.0-1.0), None if unavailable
+    cached_tokens_device: int = 0
+    cached_tokens_host: int = 0
+    cached_tokens_storage: int = 0
+    cached_storage_backend: str = ""
     completed_at: float = 0.0  # Unix timestamp when step completed
 
 
@@ -74,6 +78,10 @@ def _ensure_csv(self) -> None:
                     "completion_tokens",
                     "cached_tokens",
                     "kv_hit_rate",
+                    "cached_tokens_device",
+                    "cached_tokens_host",
+                    "cached_tokens_storage",
+                    "cached_storage_backend",
                     "completed_at",
                 ])
         self._csv_initialized = True
@@ -100,6 +108,10 @@ def _write_step_to_csv(self, metrics: StepMetrics) -> None:
                 completion_tokens_val,
                 cached_tokens_val,
                 kv_hit_val,
+                metrics.cached_tokens_device,
+                metrics.cached_tokens_host,
+                metrics.cached_tokens_storage,
+                metrics.cached_storage_backend,
                 round(metrics.completed_at, 4),
             ])
 
@@ -159,14 +171,16 @@ def on_request_end(
         prompt_tokens: Optional[int] = None,
         completion_tokens: Optional[int] = None,
         cached_tokens: Optional[int] = None,
+        cached_tokens_details: Optional[Dict[str, int | str]] = None,
     ) -> None:
         """Called when request completes.
-        
+
         Args:
             prompt_tokens: Number of prompt tokens from usage.prompt_tokens
             completion_tokens: Number of completion tokens from usage.completion_tokens
             cached_tokens: Number of cached tokens from usage.prompt_tokens_details.cached_tokens
                           (None if not available or null)
+            cached_tokens_details: Source decomposition from sglext.cached_tokens_details
         """
         now = time.time()
 
@@ -183,7 +197,13 @@ def on_request_end(
         kv_hit_rate: Optional[float] = None
         if prompt_tokens is not None and cached_tokens is not None and prompt_tokens > 0:
             kv_hit_rate = cached_tokens / prompt_tokens
-
+
+        source_details = cached_tokens_details or {}
+        cached_device = int(source_details.get("device", 0) or 0)
+        cached_host = int(source_details.get("host", 0) or 0)
+        cached_storage = int(source_details.get("storage", 0) or 0)
+        cached_storage_backend = str(source_details.get("storage_backend", "") or "")
+
         # Create step metrics
         step_id = len(self.step_metrics) + 1
         metrics = StepMetrics(
@@ -197,6 +217,10 @@ def on_request_end(
             completion_tokens=completion_tokens,
             cached_tokens=cached_tokens,
             kv_hit_rate=kv_hit_rate,
+            cached_tokens_device=cached_device,
+            cached_tokens_host=cached_host,
+            cached_tokens_storage=cached_storage,
+            cached_storage_backend=cached_storage_backend,
             completed_at=now,
         )
         self.step_metrics.append(metrics)

ThunderAgent/scheduler/vllm_request_processor.py

@@ -12,27 +12,37 @@
 import httpx
 from fastapi.responses import Response, StreamingResponse
 
+from ..config import get_config
+
 
 def extract_usage_info(
     payload: Any,
-) -> Tuple[Optional[int], Optional[int], Optional[int], Optional[int]]:
+) -> Tuple[
+    Optional[int],
+    Optional[int],
+    Optional[int],
+    Optional[int],
+    Dict[str, int | str],
+]:
     """Extract usage info from a vLLM response.
 
     Returns:
-        (total_tokens, prompt_tokens, completion_tokens, cached_tokens)
+        (total_tokens, prompt_tokens, completion_tokens, cached_tokens, cached_tokens_details)
         - cached_tokens is None if prompt_tokens_details is null or missing
+        - cached_tokens_details is {} if sglext.cached_tokens_details is missing
         - All None if usage is not available
     """
     if not isinstance(payload, dict):
-        return None, None, None, None
+        return None, None, None, None, {}
     usage = payload.get("usage")
     if not isinstance(usage, dict):
-        return None, None, None, None
+        return None, None, None, None, {}
 
     total_tokens = None
     prompt_tokens = None
     completion_tokens = None
     cached_tokens = None
+    cached_tokens_details: Dict[str, int | str] = {}
 
     if "total_tokens" in usage:
         val = usage.get("total_tokens")
@@ -70,7 +80,19 @@ def extract_usage_info(
     ):
         total_tokens = prompt_tokens + completion_tokens
 
-    return total_tokens, prompt_tokens, completion_tokens, cached_tokens
+    sglext = payload.get("sglext")
+    if isinstance(sglext, dict):
+        details = sglext.get("cached_tokens_details")
+        if isinstance(details, dict):
+            for key in ("device", "host", "storage"):
+                val = details.get(key)
+                if isinstance(val, (int, float)) and math.isfinite(val):
+                    cached_tokens_details[key] = int(val)
+            storage_backend = details.get("storage_backend")
+            if isinstance(storage_backend, str):
+                cached_tokens_details["storage_backend"] = storage_backend
+
+    return total_tokens, prompt_tokens, completion_tokens, cached_tokens, cached_tokens_details
 
 
 def filtered_headers(headers: httpx.Headers) -> Dict[str, str]:
@@ -94,6 +116,14 @@ def remove_program_id(payload: Dict[str, Any]) -> Dict[str, Any]:
     return payload
 
 
+def maybe_enable_cached_token_details(payload: Dict[str, Any]) -> Dict[str, Any]:
+    """Enable cache-source detail reporting when ThunderAgent profiling is enabled."""
+    payload = payload.copy()
+    if get_config().profile_enabled:
+        payload["return_cached_tokens_details"] = True
+    return payload
+
+
 # Default interval for token progress updates during streaming
 DEFAULT_TOKEN_PROGRESS_INTERVAL = 20
 
@@ -103,7 +133,10 @@ async def forward_streaming_request(
     url: str,
     payload: Dict[str, Any],
     *,
-    on_usage: Callable[[int, Optional[int], Optional[int], Optional[int]], Awaitable[None]] | None = None,
+    on_usage: Callable[
+        [int, Optional[int], Optional[int], Optional[int], Dict[str, int | str]],
+        Awaitable[None],
+    ] | None = None,
     on_first_token: Callable[[], None] | None = None,
     on_token: Callable[[], None] | None = None,
     on_token_progress: Callable[[int], None] | None = None,
@@ -122,7 +155,9 @@ async def forward_streaming_request(
         client: httpx AsyncClient to use
         url: vLLM endpoint URL
         payload: Request payload (program_id will be removed)
-        on_usage: Called with (total_tokens, prompt_tokens, completion_tokens, cached_tokens) when stream ends
+        on_usage: Called with
+            (total_tokens, prompt_tokens, completion_tokens, cached_tokens, cached_tokens_details)
+            when stream ends
         on_first_token: Called when first token is received
         on_token: Called for each token received
         on_token_progress: Called with cumulative token count at regular intervals
@@ -133,6 +168,7 @@ async def forward_streaming_request(
     """
     # Remove program_id before forwarding
     payload = remove_program_id(payload)
+    payload = maybe_enable_cached_token_details(payload)
 
     # Add stream_options to get usage info in streaming response
     if on_usage is not None:
@@ -155,6 +191,7 @@ async def iterator():
         prompt_tokens: Optional[int] = None
         completion_tokens: Optional[int] = None
         cached_tokens: Optional[int] = None
+        cached_tokens_details: Dict[str, int | str] = {}
         first_token_seen = False
         token_count = 0  # Track cumulative generated tokens
         last_reported_count = 0  # Last count reported via on_token_progress
@@ -194,20 +231,27 @@ async def iterator():
                             payload_obj = json.loads(data)
                         except Exception:
                             continue
-                        tt, pt, comp_t, ct = extract_usage_info(payload_obj)
+                        tt, pt, comp_t, ct, ctd = extract_usage_info(payload_obj)
                         if tt is not None:
                             total_tokens = tt
                             prompt_tokens = pt
                             completion_tokens = comp_t
                             cached_tokens = ct
+                            cached_tokens_details = ctd
                             usage_extracted = True
                 yield chunk
         finally:
             await resp_cm.__aexit__(None, None, None)
             # We assume successful vLLM/SGLang responses include usage.total_tokens
             # when include_usage is requested; otherwise this step is not finalized.
             if total_tokens is not None and on_usage is not None:
-                await on_usage(total_tokens, prompt_tokens, completion_tokens, cached_tokens)
+                await on_usage(
+                    total_tokens,
+                    prompt_tokens,
+                    completion_tokens,
+                    cached_tokens,
+                    cached_tokens_details,
+                )
 
     return StreamingResponse(
         iterator(),
@@ -222,21 +266,27 @@ async def forward_non_streaming_request(
     url: str,
     payload: Dict[str, Any],
     *,
-    on_usage: Callable[[int, Optional[int], Optional[int], Optional[int]], Awaitable[None]] | None = None,
+    on_usage: Callable[
+        [int, Optional[int], Optional[int], Optional[int], Dict[str, int | str]],
+        Awaitable[None],
+    ] | None = None,
 ) -> Response:
     """Forward a non-streaming request to vLLM and return a Response.
 
     Args:
         client: httpx AsyncClient to use
         url: vLLM endpoint URL
         payload: Request payload (program_id will be removed)
-        on_usage: Called with (total_tokens, prompt_tokens, completion_tokens, cached_tokens) after response
+        on_usage: Called with
+            (total_tokens, prompt_tokens, completion_tokens, cached_tokens, cached_tokens_details)
+            after response
 
     Returns:
         FastAPI Response with vLLM response content
     """
     # Remove program_id before forwarding
     payload = remove_program_id(payload)
+    payload = maybe_enable_cached_token_details(payload)
 
     resp = await client.post(url, json=payload)
 
@@ -245,21 +295,29 @@ async def forward_non_streaming_request(
     prompt_tokens: Optional[int] = None
     completion_tokens: Optional[int] = None
     cached_tokens: Optional[int] = None
+    cached_tokens_details: Dict[str, int | str] = {}
     try:
         payload_obj = resp.json()
     except Exception:
         payload_obj = None
-    tt, pt, comp_t, ct = extract_usage_info(payload_obj)
+    tt, pt, comp_t, ct, ctd = extract_usage_info(payload_obj)
     if tt is not None:
         total_tokens = tt
         prompt_tokens = pt
         completion_tokens = comp_t
         cached_tokens = ct
+        cached_tokens_details = ctd
 
     # Call usage callback
     # We assume successful vLLM/SGLang responses include usage.total_tokens.
     if total_tokens is not None and on_usage is not None:
-        await on_usage(total_tokens, prompt_tokens, completion_tokens, cached_tokens)
+        await on_usage(
+            total_tokens,
+            prompt_tokens,
+            completion_tokens,
+            cached_tokens,
+            cached_tokens_details,
+        )
 
     return Response(
         content=resp.content,