From 812c3d1cf30205a8c2f98ef3db7754444a9b371c Mon Sep 17 00:00:00 2001 From: AMD-yanfeiwang Date: Tue, 12 May 2026 13:54:28 +0000 Subject: [PATCH 1/4] metrics: capture API cache usage per turn --- trace_replay_tester.py | 218 ++++++++++++++++++++++++++++++++++++++++- 1 file changed, 214 insertions(+), 4 deletions(-) diff --git a/trace_replay_tester.py b/trace_replay_tester.py index 72d0833..31be463 100644 --- a/trace_replay_tester.py +++ b/trace_replay_tester.py @@ -46,7 +46,7 @@ def filter(self, record): from dataclasses import dataclass, asdict, field from datetime import datetime, timezone from pathlib import Path -from typing import List, Dict, Optional, Tuple, Literal, Set +from typing import Any, List, Dict, Optional, Tuple, Literal, Set from collections import defaultdict, deque import numpy as np import pandas as pd @@ -257,6 +257,14 @@ class RequestMetrics: # Token chunk timing for proportional output attribution token_timestamps: List[float] = field(default_factory=list) tokens_per_chunk: List[int] = field(default_factory=list) + # Actual per-request KV/prefix cache stats when exposed by the OpenAI-compatible API. + # Current SGLang responses may leave these unavailable; the fields are kept + # nullable so downstream analysis can distinguish "not reported" from 0 hits. + actual_kv_cache_hit_tokens: Optional[int] = None + actual_kv_cache_miss_tokens: Optional[int] = None + actual_kv_cache_query_tokens: Optional[int] = None + actual_kv_cache_hit_rate: Optional[float] = None + actual_kv_cache_source: Optional[str] = None def to_dict(self) -> dict: return asdict(self) @@ -1647,6 +1655,124 @@ def delta_field_names_for(model: str) -> Tuple[str, str]: return _DEFAULT_DELTA_FIELDS +def _usage_to_dict(usage: Any) -> Dict[str, Any]: + """Convert OpenAI SDK usage objects to a plain dict, preserving extras.""" + if usage is None: + return {} + if isinstance(usage, dict): + return usage + if hasattr(usage, "model_dump"): + try: + return usage.model_dump() + except Exception: + pass + result = {} + for name in dir(usage): + if name.startswith("_"): + continue + try: + value = getattr(usage, name) + except Exception: + continue + if callable(value): + continue + if isinstance(value, (str, int, float, bool, dict, list, tuple, type(None))): + result[name] = value + return result + + +def _to_optional_int(value: Any) -> Optional[int]: + if value is None: + return None + try: + return int(float(value)) + except (TypeError, ValueError): + return None + + +def _nested_get(data: Dict[str, Any], path: str) -> Any: + cur: Any = data + for part in path.split("."): + if not isinstance(cur, dict) or part not in cur: + return None + cur = cur[part] + return cur + + +def extract_api_cache_stats(usage: Any) -> Dict[str, Any]: + """Extract per-request cache stats from OpenAI-compatible usage metadata. + + Providers that expose automatic prefix-cache accounting usually place the + cached prompt token count in one of these fields: + - prompt_tokens_details.cached_tokens + - input_tokens_details.cached_tokens + - cached_tokens + + SGLang currently returns prompt_tokens_details=null on this deployment, so + this helper will mark stats unavailable until the server exposes such data. + """ + usage_dict = _usage_to_dict(usage) + prompt_tokens = _to_optional_int( + usage_dict.get("prompt_tokens", usage_dict.get("input_tokens")) + ) + + hit_candidates = [ + "prompt_tokens_details.cached_tokens", + "input_tokens_details.cached_tokens", + "cached_tokens", + "cache_read_input_tokens", + "prompt_cache_hit_tokens", + "prefix_cache_hit_tokens", + "cache_hit_tokens", + ] + miss_candidates = [ + "prompt_tokens_details.uncached_tokens", + "input_tokens_details.uncached_tokens", + "cache_creation_input_tokens", + "prompt_cache_miss_tokens", + "prefix_cache_miss_tokens", + "cache_miss_tokens", + ] + + hit_tokens = None + hit_source = None + for key in hit_candidates: + value = _nested_get(usage_dict, key) if "." in key else usage_dict.get(key) + hit_tokens = _to_optional_int(value) + if hit_tokens is not None: + hit_source = key + break + + miss_tokens = None + for key in miss_candidates: + value = _nested_get(usage_dict, key) if "." in key else usage_dict.get(key) + miss_tokens = _to_optional_int(value) + if miss_tokens is not None: + break + + query_tokens = None + if hit_tokens is not None: + explicit_total = hit_tokens + (miss_tokens or 0) + if prompt_tokens is not None: + query_tokens = max(prompt_tokens, explicit_total) + elif explicit_total > 0: + query_tokens = explicit_total + + hit_rate = None + if hit_tokens is not None and query_tokens and query_tokens > 0: + hit_rate = hit_tokens / query_tokens + + return { + "usage": usage_dict, + "hit_tokens": hit_tokens, + "miss_tokens": miss_tokens, + "query_tokens": query_tokens, + "hit_rate": hit_rate, + "source": hit_source, + "available": hit_tokens is not None and query_tokens is not None, + } + + class APIClient: """Manages OpenAI API client""" @@ -1755,6 +1881,8 @@ async def send_request(self, messages: List[dict], max_tokens: Optional[int], st token_timestamps: List[float] = [] tokens_per_chunk: List[int] = [] server_prompt_tokens: Optional[int] = None + server_usage = None + server_cache_stats: Dict[str, Any] = extract_api_cache_stats(None) content_field, reasoning_field = delta_field_names_for(self.model) debug_chunks = [] if self.debug_trace else None completion_token_ids: Optional[List[int]] = [] if self.debug_trace else None @@ -1786,7 +1914,11 @@ async def send_request(self, messages: List[dict], max_tokens: Optional[int], st except Exception: debug_chunks.append({'_repr': repr(chunk)}) if getattr(chunk, 'usage', None) is not None: - server_prompt_tokens = chunk.usage.prompt_tokens + server_usage = chunk.usage + server_cache_stats = extract_api_cache_stats(chunk.usage) + server_prompt_tokens = _to_optional_int( + server_cache_stats.get('usage', {}).get('prompt_tokens') + ) if completion_token_ids is not None and chunk.choices: chunk_logprobs = getattr(chunk.choices[0], 'logprobs', None) if chunk_logprobs is not None and tokenizer is not None: @@ -1850,7 +1982,11 @@ async def send_request(self, messages: List[dict], max_tokens: Optional[int], st response_text += reasoning_text_full token_count = response.usage.completion_tokens if response.usage else len(response_text.split()) if response.usage: - server_prompt_tokens = response.usage.prompt_tokens + server_usage = response.usage + server_cache_stats = extract_api_cache_stats(response.usage) + server_prompt_tokens = _to_optional_int( + server_cache_stats.get('usage', {}).get('prompt_tokens') + ) # For non-streaming, all tokens are attributed to completion time token_timestamps.append(complete_time) tokens_per_chunk.append(token_count) @@ -1864,6 +2000,12 @@ async def send_request(self, messages: List[dict], max_tokens: Optional[int], st 'ttlt': ttlt, 'output_tokens': token_count, 'server_prompt_tokens': server_prompt_tokens, + 'server_usage': server_cache_stats.get('usage') or _usage_to_dict(server_usage), + 'server_cache_hit_tokens': server_cache_stats.get('hit_tokens'), + 'server_cache_miss_tokens': server_cache_stats.get('miss_tokens'), + 'server_cache_query_tokens': server_cache_stats.get('query_tokens'), + 'server_cache_hit_rate': server_cache_stats.get('hit_rate'), + 'server_cache_source': server_cache_stats.get('source'), 'error_type': None, 'start_time': start_time, 'first_token_time': first_token_time or start_time, @@ -1889,6 +2031,14 @@ async def send_request(self, messages: List[dict], max_tokens: Optional[int], st 'output_tokens': token_count, 'ttft': ttft, 'ttlt': ttlt, + 'usage': server_cache_stats.get('usage') or _usage_to_dict(server_usage), + 'cache_stats': { + 'hit_tokens': server_cache_stats.get('hit_tokens'), + 'miss_tokens': server_cache_stats.get('miss_tokens'), + 'query_tokens': server_cache_stats.get('query_tokens'), + 'hit_rate': server_cache_stats.get('hit_rate'), + 'source': server_cache_stats.get('source'), + }, 'completion_token_ids': completion_token_ids, 'completion_token_count': len(completion_token_ids) if completion_token_ids is not None else None, }, @@ -1932,6 +2082,12 @@ async def send_request(self, messages: List[dict], max_tokens: Optional[int], st 'ttlt': 0, 'output_tokens': 0, 'server_prompt_tokens': None, + 'server_usage': {}, + 'server_cache_hit_tokens': None, + 'server_cache_miss_tokens': None, + 'server_cache_query_tokens': None, + 'server_cache_hit_rate': None, + 'server_cache_source': None, 'error_type': error_type, 'start_time': start_time, 'first_token_time': start_time, @@ -2672,6 +2828,11 @@ def on_chunk(chunk_time, chunk_tokens): ttlt = result['ttlt'] actual_output = result['output_tokens'] server_prompt_tokens = result.get('server_prompt_tokens') + actual_kv_cache_hit_tokens = result.get('server_cache_hit_tokens') + actual_kv_cache_miss_tokens = result.get('server_cache_miss_tokens') + actual_kv_cache_query_tokens = result.get('server_cache_query_tokens') + actual_kv_cache_hit_rate = result.get('server_cache_hit_rate') + actual_kv_cache_source = result.get('server_cache_source') error_type = result['error_type'] request_start_time = result['start_time'] prefill_complete_time = result['first_token_time'] @@ -2750,7 +2911,12 @@ def on_chunk(chunk_time, chunk_tokens): prefill_complete_time=prefill_complete_time, request_complete_time=request_complete_time, token_timestamps=token_timestamps, - tokens_per_chunk=tokens_per_chunk + tokens_per_chunk=tokens_per_chunk, + actual_kv_cache_hit_tokens=actual_kv_cache_hit_tokens, + actual_kv_cache_miss_tokens=actual_kv_cache_miss_tokens, + actual_kv_cache_query_tokens=actual_kv_cache_query_tokens, + actual_kv_cache_hit_rate=actual_kv_cache_hit_rate, + actual_kv_cache_source=actual_kv_cache_source, ) if record_user_metrics: @@ -3286,6 +3452,50 @@ def save_results(orchestrator: TestOrchestrator, config: TestConfig): df.to_csv(output_path / "detailed_results.csv", index=False) logger.info(f"Saved detailed results: {len(df)} requests") + # Save a narrow per-turn KV cache CSV for downstream analysis. These + # columns combine the existing trace-hash theoretical hit model with + # actual per-request cache fields when the OpenAI-compatible server + # exposes them (for example prompt_tokens_details.cached_tokens). + cache_rows = [] + for m in orchestrator.all_metrics: + theoretical_total_blocks = m.cache_hit_blocks + m.cache_miss_blocks + actual_available = ( + m.actual_kv_cache_hit_tokens is not None + and m.actual_kv_cache_query_tokens is not None + and m.actual_kv_cache_query_tokens > 0 + ) + cache_rows.append({ + "user_id": m.user_id, + "request_idx": m.request_idx, + "trace_id": m.trace_id, + "timestamp": m.timestamp, + "request_start_time": m.request_start_time, + "prefill_complete_time": m.prefill_complete_time, + "request_complete_time": m.request_complete_time, + "success": m.success, + "input_tokens": m.input_tokens, + "output_tokens_actual": m.output_tokens_actual, + "theoretical_cache_hit_blocks": m.cache_hit_blocks, + "theoretical_cache_miss_blocks": m.cache_miss_blocks, + "theoretical_cache_hit_rate": ( + m.cache_hit_blocks / theoretical_total_blocks + if theoretical_total_blocks > 0 else 0.0 + ), + "actual_kv_cache_available": actual_available, + "actual_kv_cache_hit_tokens": m.actual_kv_cache_hit_tokens, + "actual_kv_cache_miss_tokens": m.actual_kv_cache_miss_tokens, + "actual_kv_cache_query_tokens": m.actual_kv_cache_query_tokens, + "actual_kv_cache_hit_rate": m.actual_kv_cache_hit_rate, + "actual_kv_cache_source": m.actual_kv_cache_source, + }) + cache_df = pd.DataFrame(cache_rows) + cache_df.to_csv(output_path / "kv_cache_turn_metrics.csv", index=False) + available_count = int(cache_df["actual_kv_cache_available"].sum()) if not cache_df.empty else 0 + logger.info( + "Saved per-turn KV cache metrics: " + f"{len(cache_df)} requests ({available_count} with API-reported actual cache stats)" + ) + # Save assessment periods if orchestrator.assessment_periods: df = pd.DataFrame([asdict(p) for p in orchestrator.assessment_periods]) From 0763f6d8a8d6db00409dbf95f2f9f5040a7b809a Mon Sep 17 00:00:00 2001 From: AMD-yanfeiwang Date: Thu, 14 May 2026 05:49:24 +0000 Subject: [PATCH 2/4] metrics: merge sglang metrics improvements --- server_metrics.py | 501 ++++++++++++++++++++++++++++++++++------- trace_replay_tester.py | 72 ++++-- 2 files changed, 473 insertions(+), 100 deletions(-) diff --git a/server_metrics.py b/server_metrics.py index 9d54c6b..0f945be 100644 --- a/server_metrics.py +++ b/server_metrics.py @@ -45,6 +45,21 @@ class MetricsSnapshot: # Prefill KV computed tokens (cumulative sum from histogram) prefill_kv_computed_tokens_sum: int = 0 prefill_kv_computed_tokens_count: int = 0 + # SGLang PD KV transfer metrics (cumulative, from histograms _sum/_count) + kv_transfer_latency_ms_sum: float = 0.0 + kv_transfer_latency_ms_count: int = 0 + kv_transfer_total_mb_sum: float = 0.0 + kv_transfer_total_mb_count: int = 0 + kv_transfer_speed_gb_s_sum: float = 0.0 + kv_transfer_speed_gb_s_count: int = 0 + # SGLang HiCache host-tier metrics + hicache_host_used_tokens: int = 0 + hicache_host_total_tokens: int = 0 + # SGLang per-tier cache hit tokens (cumulative counters) + cache_hit_tokens_l1: int = 0 # GPU/device cache + cache_hit_tokens_l2: int = 0 # host/DRAM cache + cache_hit_tokens_l3: int = 0 # storage (UMBP/SSD) + cache_miss_tokens: int = 0 # ============================================================================= @@ -57,6 +72,16 @@ def _get_value(text: str, pattern: str, default: float = 0.0) -> float: return float(match.group(1)) if match else default +def _get_values(text: str, pattern: str) -> list[float]: + """Extract all values matching a Prometheus text regex.""" + return [float(match.group(1)) for match in re.finditer(pattern, text)] + + +def _sum_values(text: str, pattern: str) -> float: + """Sum all values matching a Prometheus text regex.""" + return sum(_get_values(text, pattern)) + + class VLLMMetricsParser: """Parse vLLM Prometheus metrics (prefix: vllm:).""" @@ -112,43 +137,133 @@ def parse(self, text: str) -> MetricsSnapshot: snapshot = MetricsSnapshot(timestamp=time.time()) g = lambda p, d=0.0: _get_value(text, p, d) + has_prefill_engine = 'engine_type="prefill"' in text + has_decode_engine = 'engine_type="decode"' in text + prefill_only = has_prefill_engine and not has_decode_engine + decode_only = has_decode_engine and not has_prefill_engine + # KV cache usage — sglang reports token_usage as a ratio (0-1) - snapshot.kv_cache_usage = g(r'sglang:token_usage\{[^}]*\}\s+([\d.e+-]+)') - # Fallback: compute from num_used_tokens / max_total_num_tokens - if snapshot.kv_cache_usage == 0.0: - used = g(r'sglang:num_used_tokens\{[^}]*\}\s+([\d.e+-]+)') - total = g(r'sglang:max_total_num_tokens\{[^}]*\}\s+([\d.e+-]+)') - if total > 0: - snapshot.kv_cache_usage = used / total + full_token_usage = g(r'sglang:full_token_usage\{[^}]*\}\s+([\d.e+-]+)') + token_usage = g(r'sglang:token_usage\{[^}]*\}\s+([\d.e+-]+)') + used = g(r'sglang:num_used_tokens\{[^}]*\}\s+([\d.e+-]+)') + total = g(r'sglang:max_total_num_tokens\{[^}]*\}\s+([\d.e+-]+)') + computed_token_usage = used / total if total > 0 else 0.0 + gpu_occupancy = g(r'sglang:gpu_kv_cache_occupancy\{[^}]*\}\s+([\d.e+-]+)') + for value in (full_token_usage, token_usage, computed_token_usage, gpu_occupancy): + if value > 0: + snapshot.kv_cache_usage = value + break + + host_used = g(r'sglang:hicache_host_used_tokens\{[^}]*\}\s+([\d.e+-]+)') + host_total = g(r'sglang:hicache_host_total_tokens\{[^}]*\}\s+([\d.e+-]+)') + snapshot.hicache_host_used_tokens = int(host_used) + snapshot.hicache_host_total_tokens = int(host_total) + if host_total > 0: + snapshot.cpu_kv_cache_usage = host_used / host_total snapshot.num_requests_running = int(g(r'sglang:num_running_reqs\{[^}]*\}\s+([\d.e+-]+)')) snapshot.num_requests_waiting = int(g(r'sglang:num_queue_reqs\{[^}]*\}\s+([\d.e+-]+)')) - snapshot.prompt_tokens = int(g(r'sglang:prompt_tokens_total\{[^}]*\}\s+([\d.e+-]+)')) - snapshot.generation_tokens = int(g(r'sglang:generation_tokens_total\{[^}]*\}\s+([\d.e+-]+)')) + raw_prompt_tokens = int(g(r'sglang:prompt_tokens_total\{[^}]*\}\s+([\d.e+-]+)')) + raw_generation_tokens = int(g(r'sglang:generation_tokens_total\{[^}]*\}\s+([\d.e+-]+)')) + snapshot.prompt_tokens = 0 if decode_only else raw_prompt_tokens + snapshot.generation_tokens = 0 if prefill_only else raw_generation_tokens # Preemptions — sglang calls them "retractions" snapshot.num_preemptions = int(g(r'sglang:num_retracted_reqs\{[^}]*\}\s+([\d.e+-]+)')) - snapshot.request_success = int(g(r'sglang:num_requests_total\{[^}]*\}\s+([\d.e+-]+)')) + snapshot.request_success = 0 if prefill_only else int(g(r'sglang:num_requests_total\{[^}]*\}\s+([\d.e+-]+)')) # Token source breakdown from realtime_tokens_total (cumulative) - snapshot.prompt_tokens_local_compute = int(g( - r'sglang:realtime_tokens_total\{[^}]*mode="prefill_compute"[^}]*\}\s+([\d.e+-]+)')) - snapshot.prompt_tokens_local_cache_hit = int(g( - r'sglang:realtime_tokens_total\{[^}]*mode="prefill_cache"[^}]*\}\s+([\d.e+-]+)')) + if not decode_only: + snapshot.prompt_tokens_local_compute = int(g( + r'sglang:realtime_tokens_total\{[^}]*mode="prefill_compute"[^}]*\}\s+([\d.e+-]+)')) + realtime_cache_hit = int(g( + r'sglang:realtime_tokens_total\{[^}]*mode="prefill_cache"[^}]*\}\s+([\d.e+-]+)')) + device_cache_hit = int(g( + r'sglang:cached_tokens_total\{[^}]*cache_source="device"[^}]*\}\s+([\d.e+-]+)')) + host_cache_hit = int(g( + r'sglang:cached_tokens_total\{[^}]*cache_source="host"[^}]*\}\s+([\d.e+-]+)')) + + if device_cache_hit > 0 or host_cache_hit > 0: + snapshot.prompt_tokens_local_cache_hit = device_cache_hit + snapshot.prompt_tokens_external_kv_transfer = host_cache_hit + else: + snapshot.prompt_tokens_local_cache_hit = realtime_cache_hit + + snapshot.cache_hit_tokens_l1 = int(_sum_values( + text, + r'sglang:cache_hit_tokens_l1_total\{[^}]*\}\s+([\d.e+-]+)', + )) + snapshot.cache_hit_tokens_l2 = int(_sum_values( + text, + r'sglang:cache_hit_tokens_l2_total\{[^}]*\}\s+([\d.e+-]+)', + )) + snapshot.cache_hit_tokens_l3 = int(_sum_values( + text, + r'sglang:cache_hit_tokens_l3_total\{[^}]*\}\s+([\d.e+-]+)', + )) + snapshot.cache_miss_tokens = int(_sum_values( + text, + r'sglang:cache_miss_tokens_total\{[^}]*\}\s+([\d.e+-]+)', + )) + if not any(( + snapshot.cache_hit_tokens_l1, + snapshot.cache_hit_tokens_l2, + snapshot.cache_hit_tokens_l3, + snapshot.cache_miss_tokens, + )): + snapshot.cache_hit_tokens_l1 = device_cache_hit + snapshot.cache_hit_tokens_l2 = host_cache_hit + snapshot.cache_miss_tokens = snapshot.prompt_tokens_local_compute # Derive cumulative hits/queries from the per-source token counters. # This is the correct cumulative cache hit ratio — unlike sglang's # instantaneous `cache_hit_rate` gauge, which is 0 during decode-only # periods and thus yielded spurious 0% hit rates when sampled at # benchmark shutdown. - snapshot.prefix_cache_hits = snapshot.prompt_tokens_local_cache_hit - snapshot.prefix_cache_queries = ( + snapshot.prefix_cache_hits = ( snapshot.prompt_tokens_local_cache_hit + + snapshot.prompt_tokens_external_kv_transfer + ) + snapshot.prefix_cache_queries = ( + snapshot.prefix_cache_hits + snapshot.prompt_tokens_local_compute ) + snapshot.kv_offload_bytes_gpu_to_cpu = _sum_values( + text, + r'sglang:evicted_bytes_total\{[^}]*\}\s+([\d.e+-]+)', + ) + snapshot.kv_offload_bytes_cpu_to_gpu = _sum_values( + text, + r'sglang:load_back_bytes_total\{[^}]*\}\s+([\d.e+-]+)', + ) + snapshot.kv_transfer_latency_ms_sum = _sum_values( + text, + r'sglang:kv_transfer_latency_ms_sum\{[^}]*\}\s+([\d.e+-]+)', + ) + snapshot.kv_transfer_latency_ms_count = int(_sum_values( + text, + r'sglang:kv_transfer_latency_ms_count\{[^}]*\}\s+([\d.e+-]+)', + )) + snapshot.kv_transfer_total_mb_sum = _sum_values( + text, + r'sglang:kv_transfer_total_mb_sum\{[^}]*\}\s+([\d.e+-]+)', + ) + snapshot.kv_transfer_total_mb_count = int(_sum_values( + text, + r'sglang:kv_transfer_total_mb_count\{[^}]*\}\s+([\d.e+-]+)', + )) + snapshot.kv_transfer_speed_gb_s_sum = _sum_values( + text, + r'sglang:kv_transfer_speed_gb_s_sum\{[^}]*\}\s+([\d.e+-]+)', + ) + snapshot.kv_transfer_speed_gb_s_count = int(_sum_values( + text, + r'sglang:kv_transfer_speed_gb_s_count\{[^}]*\}\s+([\d.e+-]+)', + )) + return snapshot @@ -170,37 +285,126 @@ def get_parser(backend: str): @dataclass class MetricsCollector: - base_url: str + base_url: str | None = None poll_interval: float = 1.0 + metrics_endpoints: str | list[str] | None = None snapshots: list[MetricsSnapshot] = field(default_factory=list) _running: bool = False _task: asyncio.Task | None = None - _parser: VLLMMetricsParser | SGLangMetricsParser | None = None - _backend: str = "" + _metrics_urls: list[str] = field(default_factory=list, init=False) + _parsers: dict[str, VLLMMetricsParser | SGLangMetricsParser] = field(default_factory=dict, init=False) + _backends: dict[str, str] = field(default_factory=dict, init=False) gpu_transfer_collector: object = None - def _parse_metrics(self, text: str) -> MetricsSnapshot: + def __post_init__(self) -> None: + self._metrics_urls = self._normalize_metrics_endpoints( + self.metrics_endpoints if self.metrics_endpoints else self.base_url + ) + + @staticmethod + def _normalize_metrics_endpoints(endpoints: str | list[str] | None) -> list[str]: + if endpoints is None: + return [] + + if isinstance(endpoints, str): + entries = endpoints.split(',') + else: + entries = endpoints + + metrics_urls = [] + for entry in entries: + endpoint = str(entry).strip().rstrip('/') + if not endpoint: + continue + if not endpoint.endswith('/metrics'): + endpoint = f"{endpoint}/metrics" + metrics_urls.append(endpoint) + return metrics_urls + + def describe_endpoints(self) -> str: + return ', '.join(self._metrics_urls) if self._metrics_urls else '' + + def _parse_metrics(self, text: str, metrics_url: str) -> MetricsSnapshot: """Parse Prometheus metrics text, auto-detecting backend on first call.""" - if self._parser is None: - self._backend = detect_backend(text) - self._parser = get_parser(self._backend) - if self._backend != 'unknown': - logger.info(f"Auto-detected metrics backend: {self._backend}") - return self._parser.parse(text) + parser = self._parsers.get(metrics_url) + if parser is None: + backend = detect_backend(text) + parser = get_parser(backend) + self._parsers[metrics_url] = parser + self._backends[metrics_url] = backend + if backend != 'unknown': + logger.info(f"Auto-detected metrics backend for {metrics_url}: {backend}") + return parser.parse(text) + + async def _fetch_snapshot( + self, + session: aiohttp.ClientSession, + metrics_url: str, + ) -> MetricsSnapshot | None: + try: + async with session.get(metrics_url, timeout=aiohttp.ClientTimeout(total=5)) as resp: + if resp.status != 200: + logger.warning(f"Metrics poll error from {metrics_url}: HTTP {resp.status}") + return None + text = await resp.text() + return self._parse_metrics(text, metrics_url) + except Exception as e: + logger.warning(f"Metrics poll error from {metrics_url}: {e}") + return None + + @staticmethod + def _combine_snapshots(snapshots: list[MetricsSnapshot]) -> MetricsSnapshot: + combined = MetricsSnapshot(timestamp=time.time()) + combined.kv_cache_usage = max(s.kv_cache_usage for s in snapshots) + combined.cpu_kv_cache_usage = max(s.cpu_kv_cache_usage for s in snapshots) + combined.num_requests_running = sum(s.num_requests_running for s in snapshots) + combined.num_requests_waiting = sum(s.num_requests_waiting for s in snapshots) + combined.prefix_cache_hits = sum(s.prefix_cache_hits for s in snapshots) + combined.prefix_cache_queries = sum(s.prefix_cache_queries for s in snapshots) + combined.cpu_prefix_cache_hits = sum(s.cpu_prefix_cache_hits for s in snapshots) + combined.cpu_prefix_cache_queries = sum(s.cpu_prefix_cache_queries for s in snapshots) + combined.prompt_tokens = sum(s.prompt_tokens for s in snapshots) + combined.generation_tokens = sum(s.generation_tokens for s in snapshots) + combined.num_preemptions = sum(s.num_preemptions for s in snapshots) + combined.request_success = sum(s.request_success for s in snapshots) + combined.kv_offload_bytes_gpu_to_cpu = sum(s.kv_offload_bytes_gpu_to_cpu for s in snapshots) + combined.kv_offload_bytes_cpu_to_gpu = sum(s.kv_offload_bytes_cpu_to_gpu for s in snapshots) + combined.kv_offload_time_gpu_to_cpu = sum(s.kv_offload_time_gpu_to_cpu for s in snapshots) + combined.kv_offload_time_cpu_to_gpu = sum(s.kv_offload_time_cpu_to_gpu for s in snapshots) + combined.prompt_tokens_local_compute = sum(s.prompt_tokens_local_compute for s in snapshots) + combined.prompt_tokens_local_cache_hit = sum(s.prompt_tokens_local_cache_hit for s in snapshots) + combined.prompt_tokens_external_kv_transfer = sum(s.prompt_tokens_external_kv_transfer for s in snapshots) + combined.prefill_kv_computed_tokens_sum = sum(s.prefill_kv_computed_tokens_sum for s in snapshots) + combined.prefill_kv_computed_tokens_count = sum(s.prefill_kv_computed_tokens_count for s in snapshots) + combined.kv_transfer_latency_ms_sum = sum(s.kv_transfer_latency_ms_sum for s in snapshots) + combined.kv_transfer_latency_ms_count = sum(s.kv_transfer_latency_ms_count for s in snapshots) + combined.kv_transfer_total_mb_sum = sum(s.kv_transfer_total_mb_sum for s in snapshots) + combined.kv_transfer_total_mb_count = sum(s.kv_transfer_total_mb_count for s in snapshots) + combined.kv_transfer_speed_gb_s_sum = sum(s.kv_transfer_speed_gb_s_sum for s in snapshots) + combined.kv_transfer_speed_gb_s_count = sum(s.kv_transfer_speed_gb_s_count for s in snapshots) + combined.hicache_host_used_tokens = max(s.hicache_host_used_tokens for s in snapshots) + combined.hicache_host_total_tokens = max(s.hicache_host_total_tokens for s in snapshots) + combined.cache_hit_tokens_l1 = sum(s.cache_hit_tokens_l1 for s in snapshots) + combined.cache_hit_tokens_l2 = sum(s.cache_hit_tokens_l2 for s in snapshots) + combined.cache_hit_tokens_l3 = sum(s.cache_hit_tokens_l3 for s in snapshots) + combined.cache_miss_tokens = sum(s.cache_miss_tokens for s in snapshots) + return combined async def _poll_loop(self) -> None: """Background polling loop.""" - metrics_url = f"{self.base_url}/metrics" + if not self._metrics_urls: + logger.warning("No metrics endpoints configured") + return + async with aiohttp.ClientSession() as session: while self._running: - try: - async with session.get(metrics_url, timeout=aiohttp.ClientTimeout(total=5)) as resp: - if resp.status == 200: - text = await resp.text() - snapshot = self._parse_metrics(text) - self.snapshots.append(snapshot) - except Exception as e: - logger.warning(f"Metrics poll error: {e}") + results = await asyncio.gather( + *(self._fetch_snapshot(session, url) for url in self._metrics_urls), + return_exceptions=True, + ) + snapshots = [s for s in results if isinstance(s, MetricsSnapshot)] + if snapshots: + self.snapshots.append(self._combine_snapshots(snapshots)) await asyncio.sleep(self.poll_interval) @@ -262,7 +466,7 @@ def generate_plots( # Create figure with subplots num_rows = 6 if client_metrics else 4 fig, axes = plt.subplots(num_rows, 2, figsize=(14, 4 * num_rows)) - fig.suptitle("vLLM Server Metrics During Benchmark", fontsize=14) + fig.suptitle("Server Metrics During Benchmark", fontsize=14) # 1. KV Cache Usage vs Time ax = axes[0, 0] @@ -437,42 +641,70 @@ def generate_plots( ax.set_title("Throughput (Total & Decode)") ax.grid(True, alpha=0.3) - # 5. KV Offload Transfer Rate (from vLLM metrics) + # 5. KV Transfer Rate ax = axes[2, 0] - gpu_to_cpu_rates = [] - cpu_to_gpu_rates = [] - for i in range(1, len(self.snapshots)): - dt = self.snapshots[i].timestamp - self.snapshots[i-1].timestamp - if dt > 0: - delta_g2c = self.snapshots[i].kv_offload_bytes_gpu_to_cpu - self.snapshots[i-1].kv_offload_bytes_gpu_to_cpu - delta_c2g = self.snapshots[i].kv_offload_bytes_cpu_to_gpu - self.snapshots[i-1].kv_offload_bytes_cpu_to_gpu - gpu_to_cpu_rates.append(delta_g2c / dt / 1e6) # MB/s - cpu_to_gpu_rates.append(delta_c2g / dt / 1e6) # MB/s - else: - gpu_to_cpu_rates.append(0) - cpu_to_gpu_rates.append(0) - if any(r > 0 for r in gpu_to_cpu_rates) or any(r > 0 for r in cpu_to_gpu_rates): - ax.scatter(times[1:], gpu_to_cpu_rates, alpha=0.15, s=3, c='blue') - ax.scatter(times[1:], cpu_to_gpu_rates, alpha=0.15, s=3, c='red') - xfer_window = min(30, len(gpu_to_cpu_rates) // 10) if len(gpu_to_cpu_rates) > 10 else 1 - if xfer_window > 1: - rolling_g2c = [ - sum(gpu_to_cpu_rates[max(0, i - xfer_window):i + 1]) / len(gpu_to_cpu_rates[max(0, i - xfer_window):i + 1]) - for i in range(len(gpu_to_cpu_rates)) - ] - rolling_c2g = [ - sum(cpu_to_gpu_rates[max(0, i - xfer_window):i + 1]) / len(cpu_to_gpu_rates[max(0, i - xfer_window):i + 1]) - for i in range(len(cpu_to_gpu_rates)) - ] - ax.plot(times[1:], rolling_g2c, 'b-', linewidth=1.5, label=f'GPU→CPU (avg n={xfer_window})') - ax.plot(times[1:], rolling_c2g, 'r-', linewidth=1.5, label=f'CPU→GPU (avg n={xfer_window})') - else: - ax.plot(times[1:], gpu_to_cpu_rates, 'b-', linewidth=1, alpha=0.8, label='GPU→CPU') - ax.plot(times[1:], cpu_to_gpu_rates, 'r-', linewidth=1, alpha=0.8, label='CPU→GPU') - ax.legend(fontsize=8) + has_pd_transfer = any(s.kv_transfer_total_mb_count > 0 for s in self.snapshots) + if has_pd_transfer: + interval_speeds = [] + interval_mb = [] + for i in range(1, len(self.snapshots)): + dc = self.snapshots[i].kv_transfer_speed_gb_s_count - self.snapshots[i-1].kv_transfer_speed_gb_s_count + ds = self.snapshots[i].kv_transfer_speed_gb_s_sum - self.snapshots[i-1].kv_transfer_speed_gb_s_sum + dm = self.snapshots[i].kv_transfer_total_mb_sum - self.snapshots[i-1].kv_transfer_total_mb_sum + interval_speeds.append(ds / dc if dc > 0 else 0) + interval_mb.append(dm) + if any(v > 0 for v in interval_speeds): + ax.scatter(times[1:], interval_speeds, alpha=0.15, s=3, c='blue') + xfer_window = max(3, min(30, len(interval_speeds) // 5)) if len(interval_speeds) >= 3 else len(interval_speeds) + if xfer_window > 1: + rolling = [ + sum(interval_speeds[max(0, i - xfer_window):i + 1]) / len(interval_speeds[max(0, i - xfer_window):i + 1]) + for i in range(len(interval_speeds)) + ] + ax.plot(times[1:], rolling, 'b-', linewidth=1.5, label=f'Rolling avg (n={xfer_window})') + ax.legend(fontsize=8) + ax2 = ax.twinx() + width = max(0.5, (times[-1] - times[0]) / max(1, len(times)) * 0.8) + ax2.bar(times[1:], interval_mb, width=width, alpha=0.15, color='orange', label='MB transferred') + ax2.set_ylabel("MB/interval", color='orange') + ax2.legend(fontsize=7, loc='upper left') + ax.set_ylabel("Avg Speed (GB/s)") + ax.set_title("KV Transfer Speed (PD Disagg)") + else: + gpu_to_cpu_rates = [] + cpu_to_gpu_rates = [] + for i in range(1, len(self.snapshots)): + dt = self.snapshots[i].timestamp - self.snapshots[i-1].timestamp + if dt > 0: + delta_g2c = self.snapshots[i].kv_offload_bytes_gpu_to_cpu - self.snapshots[i-1].kv_offload_bytes_gpu_to_cpu + delta_c2g = self.snapshots[i].kv_offload_bytes_cpu_to_gpu - self.snapshots[i-1].kv_offload_bytes_cpu_to_gpu + gpu_to_cpu_rates.append(delta_g2c / dt / 1e6) # MB/s + cpu_to_gpu_rates.append(delta_c2g / dt / 1e6) # MB/s + else: + gpu_to_cpu_rates.append(0) + cpu_to_gpu_rates.append(0) + if any(r > 0 for r in gpu_to_cpu_rates) or any(r > 0 for r in cpu_to_gpu_rates): + ax.scatter(times[1:], gpu_to_cpu_rates, alpha=0.15, s=3, c='blue') + ax.scatter(times[1:], cpu_to_gpu_rates, alpha=0.15, s=3, c='red') + xfer_window = min(30, len(gpu_to_cpu_rates) // 10) if len(gpu_to_cpu_rates) > 10 else 1 + if xfer_window > 1: + rolling_g2c = [ + sum(gpu_to_cpu_rates[max(0, i - xfer_window):i + 1]) / len(gpu_to_cpu_rates[max(0, i - xfer_window):i + 1]) + for i in range(len(gpu_to_cpu_rates)) + ] + rolling_c2g = [ + sum(cpu_to_gpu_rates[max(0, i - xfer_window):i + 1]) / len(cpu_to_gpu_rates[max(0, i - xfer_window):i + 1]) + for i in range(len(cpu_to_gpu_rates)) + ] + ax.plot(times[1:], rolling_g2c, 'b-', linewidth=1.5, label=f'GPU→CPU (avg n={xfer_window})') + ax.plot(times[1:], rolling_c2g, 'r-', linewidth=1.5, label=f'CPU→GPU (avg n={xfer_window})') + else: + ax.plot(times[1:], gpu_to_cpu_rates, 'b-', linewidth=1, alpha=0.8, label='GPU→CPU') + ax.plot(times[1:], cpu_to_gpu_rates, 'r-', linewidth=1, alpha=0.8, label='CPU→GPU') + ax.legend(fontsize=8) + ax.set_ylabel("Transfer Rate (MB/s)") + ax.set_title("KV Offload Transfer Rate") ax.set_xlabel("Time (s)") - ax.set_ylabel("Transfer Rate (MB/s)") - ax.set_title("KV Offload Transfer Rate") ax.grid(True, alpha=0.3) # 6. Prompt Token Sources Over Time (cumulative percentage) @@ -505,30 +737,69 @@ def generate_plots( ax.set_ylim(0, 105) ax.grid(True, alpha=0.3) - # 7. Cumulative KV Offload Transfers + # 7. HiCache / Cumulative KV Offload Transfers initial = self.snapshots[0] - # GPU → CPU cumulative + has_hicache = any(s.hicache_host_total_tokens > 0 for s in self.snapshots) + has_tier_cache = any( + s.cache_hit_tokens_l1 > 0 + or s.cache_hit_tokens_l2 > 0 + or s.cache_hit_tokens_l3 > 0 + or s.cache_miss_tokens > 0 + for s in self.snapshots + ) + ax = axes[3, 0] - cum_g2c = [(s.kv_offload_bytes_gpu_to_cpu - initial.kv_offload_bytes_gpu_to_cpu) / 1e9 - for s in self.snapshots] - if any(v > 0 for v in cum_g2c): - ax.plot(times, cum_g2c, 'b-', linewidth=1.5) - ax.fill_between(times, cum_g2c, alpha=0.2, color='blue') + if has_hicache: + host_usage = [ + 100.0 * s.hicache_host_used_tokens / s.hicache_host_total_tokens + if s.hicache_host_total_tokens > 0 else 0.0 + for s in self.snapshots + ] + host_used = [s.hicache_host_used_tokens / 1e6 for s in self.snapshots] + ax.plot(times, host_usage, 'g-', linewidth=1.5, label='Host tier usage') + ax.fill_between(times, host_usage, alpha=0.15, color='green') + ax2 = ax.twinx() + ax2.plot(times, host_used, 'b--', linewidth=1, alpha=0.7, label='Used tokens') + ax2.set_ylabel("Used Tokens (M)", color='blue') + ax2.tick_params(axis='y', labelcolor='blue') + ax.set_ylabel("Host KV Usage (%)", color='green') + ax.tick_params(axis='y', labelcolor='green') + ax.set_title("HiCache Host-Tier Usage") + else: + cum_g2c = [(s.kv_offload_bytes_gpu_to_cpu - initial.kv_offload_bytes_gpu_to_cpu) / 1e9 + for s in self.snapshots] + if any(v > 0 for v in cum_g2c): + ax.plot(times, cum_g2c, 'b-', linewidth=1.5) + ax.fill_between(times, cum_g2c, alpha=0.2, color='blue') + ax.set_ylabel("Cumulative Transfer (GB)") + ax.set_title("KV Offload: GPU → CPU (Cumulative)") ax.set_xlabel("Time (s)") - ax.set_ylabel("Cumulative Transfer (GB)") - ax.set_title("KV Offload: GPU → CPU (Cumulative)") ax.grid(True, alpha=0.3) - # CPU → GPU cumulative ax = axes[3, 1] - cum_c2g = [(s.kv_offload_bytes_cpu_to_gpu - initial.kv_offload_bytes_cpu_to_gpu) / 1e9 - for s in self.snapshots] - if any(v > 0 for v in cum_c2g): - ax.plot(times, cum_c2g, 'r-', linewidth=1.5) - ax.fill_between(times, cum_c2g, alpha=0.2, color='red') + if has_tier_cache: + l1 = [max(0, s.cache_hit_tokens_l1 - initial.cache_hit_tokens_l1) for s in self.snapshots] + l2 = [max(0, s.cache_hit_tokens_l2 - initial.cache_hit_tokens_l2) for s in self.snapshots] + l3 = [max(0, s.cache_hit_tokens_l3 - initial.cache_hit_tokens_l3) for s in self.snapshots] + miss = [max(0, s.cache_miss_tokens - initial.cache_miss_tokens) for s in self.snapshots] + ax.stackplot( + times, l1, l2, l3, miss, + labels=['L1 GPU', 'L2 Host', 'L3 Storage', 'Miss'], + colors=['steelblue', 'mediumseagreen', 'gold', 'coral'], + alpha=0.75, + ) + ax.legend(fontsize=8, loc='upper left') + ax.set_ylabel("Cumulative Tokens") + ax.set_title("Cache Hit/Miss Breakdown by Tier") + else: + cum_c2g = [(s.kv_offload_bytes_cpu_to_gpu - initial.kv_offload_bytes_cpu_to_gpu) / 1e9 + for s in self.snapshots] + if any(v > 0 for v in cum_c2g): + ax.plot(times, cum_c2g, 'r-', linewidth=1.5) + ax.fill_between(times, cum_c2g, alpha=0.2, color='red') + ax.set_ylabel("Cumulative Transfer (GB)") + ax.set_title("KV Offload: CPU → GPU (Cumulative)") ax.set_xlabel("Time (s)") - ax.set_ylabel("Cumulative Transfer (GB)") - ax.set_title("KV Offload: CPU → GPU (Cumulative)") ax.grid(True, alpha=0.3) # 8 & 9. Client metrics plots (TTFT and Latency vs Time) @@ -691,6 +962,36 @@ def _print_summary(self) -> None: lines.append(f" GPU→CPU: {g2c_bytes/1e9:.2f} GB in {g2c_time:.2f}s ({g2c_bytes/g2c_time/1e9:.1f} GB/s)" if g2c_time > 0 else f" GPU→CPU: {g2c_bytes/1e9:.2f} GB") lines.append(f" CPU→GPU: {c2g_bytes/1e9:.2f} GB in {c2g_time:.2f}s ({c2g_bytes/c2g_time/1e9:.1f} GB/s)" if c2g_time > 0 else f" CPU→GPU: {c2g_bytes/1e9:.2f} GB") + kv_transfer_count = final.kv_transfer_total_mb_count - initial.kv_transfer_total_mb_count + if kv_transfer_count > 0: + kv_transfer_mb = final.kv_transfer_total_mb_sum - initial.kv_transfer_total_mb_sum + speed_count = final.kv_transfer_speed_gb_s_count - initial.kv_transfer_speed_gb_s_count + speed_sum = final.kv_transfer_speed_gb_s_sum - initial.kv_transfer_speed_gb_s_sum + latency_count = final.kv_transfer_latency_ms_count - initial.kv_transfer_latency_ms_count + latency_sum = final.kv_transfer_latency_ms_sum - initial.kv_transfer_latency_ms_sum + lines.append(f"SGLang PD KV transfers:") + lines.append(f" Total: {kv_transfer_mb:.1f} MB across {kv_transfer_count:,} transfers") + if speed_count > 0: + lines.append(f" Avg speed: {speed_sum/speed_count:.2f} GB/s") + if latency_count > 0: + lines.append(f" Avg latency: {latency_sum/latency_count:.2f} ms") + + if final.hicache_host_total_tokens > 0: + host_usage = 100.0 * final.hicache_host_used_tokens / final.hicache_host_total_tokens + lines.append(f"HiCache host tier: {final.hicache_host_used_tokens:,}/{final.hicache_host_total_tokens:,} tokens ({host_usage:.1f}%)") + + delta_l1 = final.cache_hit_tokens_l1 - initial.cache_hit_tokens_l1 + delta_l2 = final.cache_hit_tokens_l2 - initial.cache_hit_tokens_l2 + delta_l3 = final.cache_hit_tokens_l3 - initial.cache_hit_tokens_l3 + delta_miss = final.cache_miss_tokens - initial.cache_miss_tokens + tier_total = delta_l1 + delta_l2 + delta_l3 + delta_miss + if tier_total > 0: + lines.append(f"Cache tier token breakdown:") + lines.append(f" - L1 GPU: {delta_l1:>12,} ({100*delta_l1/tier_total:.1f}%)") + lines.append(f" - L2 Host: {delta_l2:>12,} ({100*delta_l2/tier_total:.1f}%)") + lines.append(f" - L3 Storage: {delta_l3:>12,} ({100*delta_l3/tier_total:.1f}%)") + lines.append(f" - Miss: {delta_miss:>12,} ({100*delta_miss/tier_total:.1f}%)") + delta_kv_sum = final.prefill_kv_computed_tokens_sum - initial.prefill_kv_computed_tokens_sum delta_kv_count = final.prefill_kv_computed_tokens_count - initial.prefill_kv_computed_tokens_count if delta_kv_count > 0: @@ -758,6 +1059,20 @@ def export_csv( # Prefill KV computed 'prefill_kv_computed_tokens_sum', 'prefill_kv_computed_tokens_count', + # SGLang PD KV transfer histograms + 'kv_transfer_latency_ms_sum', + 'kv_transfer_latency_ms_count', + 'kv_transfer_total_mb_sum', + 'kv_transfer_total_mb_count', + 'kv_transfer_speed_gb_s_sum', + 'kv_transfer_speed_gb_s_count', + # SGLang HiCache / cache tiers + 'hicache_host_used_tokens', + 'hicache_host_total_tokens', + 'cache_hit_tokens_l1', + 'cache_hit_tokens_l2', + 'cache_hit_tokens_l3', + 'cache_miss_tokens', # Computed per-interval metrics 'interval_cache_hit_rate_pct', 'interval_throughput_tok_per_sec', @@ -805,6 +1120,18 @@ def export_csv( s.prompt_tokens_external_kv_transfer, s.prefill_kv_computed_tokens_sum, s.prefill_kv_computed_tokens_count, + f"{s.kv_transfer_latency_ms_sum:.6f}", + s.kv_transfer_latency_ms_count, + f"{s.kv_transfer_total_mb_sum:.6f}", + s.kv_transfer_total_mb_count, + f"{s.kv_transfer_speed_gb_s_sum:.6f}", + s.kv_transfer_speed_gb_s_count, + s.hicache_host_used_tokens, + s.hicache_host_total_tokens, + s.cache_hit_tokens_l1, + s.cache_hit_tokens_l2, + s.cache_hit_tokens_l3, + s.cache_miss_tokens, f"{cache_hit_rate:.2f}", f"{throughput:.2f}", ]) diff --git a/trace_replay_tester.py b/trace_replay_tester.py index 31be463..67852f1 100644 --- a/trace_replay_tester.py +++ b/trace_replay_tester.py @@ -203,8 +203,9 @@ class TestConfig: tokenizer_id: str min_requests: int = 1 seed: Optional[int] = None # Random seed for reproducibility - # Generation parameters (None = use model defaults or auto-detect) - temperature: Optional[float] = None + # Generation parameters. Default to greedy sampling to avoid top-p/top-k + # filtering effects during stability/debug runs. + temperature: Optional[float] = 0.0 top_p: Optional[float] = None top_k: Optional[int] = None repetition_penalty: Optional[float] = None @@ -222,7 +223,9 @@ class TestConfig: hash_block_mode: bool = False debug_trace: bool = False # Store full request/response bodies to JSONL metrics_output_prefix: Optional[str] = None # Enable server metrics collection + metrics_endpoint: Optional[str] = None # Comma-separated /metrics endpoints or base URLs warmup_enabled: bool = False # Pre-send one request per advanced user before metrics start + request_rate: float = float("inf") # Initial dispatch Poisson rate; inf = dispatch immediately def to_dict(self) -> dict: return asdict(self) @@ -1936,7 +1939,7 @@ async def send_request(self, messages: List[dict], max_tokens: Optional[int], st delta = chunk.choices[0].delta content_text = getattr(delta, content_field, None) or "" reasoning_text = getattr(delta, reasoning_field, None) or "" - chunk_text = content_text or reasoning_text + chunk_text = reasoning_text + content_text if not chunk_text: continue chunk_time = time.time() @@ -1951,11 +1954,11 @@ async def send_request(self, messages: List[dict], max_tokens: Optional[int], st # reasoning + content + harmony framing — so to keep the # next turn's prompt size faithful to the trace we have to # carry whatever channel(s) were actually used. - if content_text: - response_text += content_text if reasoning_text: response_text += reasoning_text reasoning_text_full += reasoning_text + if content_text: + response_text += content_text # Count all generated tokens (content + reasoning) for metrics chunk_token_count = len(tokenizer.encode(chunk_text)) if tokenizer else 1 token_count += chunk_token_count @@ -1976,10 +1979,9 @@ async def send_request(self, messages: List[dict], max_tokens: Optional[int], st if response.choices: msg = response.choices[0].message - response_text = getattr(msg, content_field, None) or "" + content_text = getattr(msg, content_field, None) or "" reasoning_text_full = getattr(msg, reasoning_field, None) or "" - if reasoning_text_full: - response_text += reasoning_text_full + response_text = reasoning_text_full + content_text token_count = response.usage.completion_tokens if response.usage else len(response_text.split()) if response.usage: server_usage = response.usage @@ -2153,6 +2155,7 @@ def __init__(self, config: TestConfig, trace_manager: TraceManager, if config.metrics_output_prefix and MetricsCollector is not None: self.metrics_collector = MetricsCollector( base_url=config.api_endpoint, + metrics_endpoints=config.metrics_endpoint, poll_interval=1.0, ) @@ -3024,11 +3027,35 @@ async def run(self): if self.metrics_collector: self.metrics_collector.start() - logger.info(f"Server metrics collector started (polling {self.config.api_endpoint}/metrics)") + logger.info(f"Server metrics collector started (polling {self.metrics_collector.describe_endpoints()})") # Track in-flight tasks: maps task -> (user_id, start_time) pending_tasks: Dict[asyncio.Task, Tuple[str, float]] = {} + if self.config.request_rate != float("inf"): + rate = self.config.request_rate + if rate <= 0: + raise ValueError("--request-rate must be positive or inf") + logger.info(f"Staggered first dispatch: Poisson rate={rate} req/s") + for user_id, user in list(self.users.items()): + if user.state != "idle": + continue + + self.process_subagent_markers(user) + if user.current_idx >= len(user.requests) and not user.pending_subagents: + user.state = "completed" + continue + if user.pending_subagents: + continue + if (self.config.max_concurrent_requests and + self.in_flight_requests >= self.config.max_concurrent_requests): + break + + self.in_flight_requests += 1 + task = asyncio.create_task(self.run_user_request(user)) + pending_tasks[task] = (user_id, time.time()) + await asyncio.sleep(float(np.random.exponential(1.0 / rate))) + try: while self.running: # Check test duration @@ -3621,10 +3648,20 @@ def parse_arguments(): "results/metrics_plots.png). Polls the server /metrics endpoint " "during the test and generates plots/CSVs with both server and " "client metrics.") - - # Generation parameter overrides (None = use model-specific defaults if available) - parser.add_argument("--temperature", type=float, default=None, - help="Override temperature for generation (e.g., 0.7)") + parser.add_argument("--metrics-endpoint", type=str, default=None, + help="Comma-separated server metrics endpoints. Each entry may be a " + "base URL (http://host:port) or a full /metrics URL. Defaults to " + "the API endpoint when omitted.") + + parser.add_argument("--request-rate", type=float, default=float("inf"), + help="Initial request dispatch rate in requests/sec. Use inf to dispatch " + "the first batch immediately; finite values use Poisson inter-arrival " + "times to avoid a synchronized first burst.") + + # Generation parameter overrides. The default temperature=0.0 selects greedy + # sampling in SGLang, which normalizes internally to top_k=1. + parser.add_argument("--temperature", type=float, default=0.0, + help="Override temperature for generation (default: 0.0 greedy; e.g., 0.7)") parser.add_argument("--top-p", type=float, default=None, help="Override top_p for generation (e.g., 0.8)") parser.add_argument("--top-k", type=int, default=None, @@ -3754,7 +3791,9 @@ async def main(): hash_block_mode=args.hash_block_mode, debug_trace=args.debug_trace, metrics_output_prefix=args.metrics_output_prefix, + metrics_endpoint=args.metrics_endpoint, warmup_enabled=args.warmup_enabled, + request_rate=args.request_rate, ) # Print header @@ -3810,6 +3849,13 @@ async def main(): logger.info(f" Test Duration: {config.test_duration}s") if config.max_concurrent_requests: logger.info(f" Max Concurrent Requests: {config.max_concurrent_requests}") + if config.request_rate != float("inf"): + logger.info(f" Initial Request Rate: {config.request_rate} req/s") + logger.info( + f" Sampling: temperature={config.temperature}, " + f"top_p={config.top_p}, top_k={config.top_k}, " + f"repetition_penalty={config.repetition_penalty}" + ) if config.warm_prefix_pct > 0: warm_tokens = int(config.warm_prefix_pct * stats.max_shared_prefix_tokens) if stats.max_shared_prefix_tokens > 0 else 0 logger.info(f" Warm Prefix: {config.warm_prefix_pct:.0%} of tool+system ({warm_tokens:,} tokens)") From c784f6a2589722d66e597c1ff0179b9b429c5d37 Mon Sep 17 00:00:00 2001 From: AMD-yanfeiwang Date: Thu, 14 May 2026 09:50:48 +0000 Subject: [PATCH 3/4] trace-replay: send mori trace headers --- trace_replay_tester.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/trace_replay_tester.py b/trace_replay_tester.py index 67852f1..21d52df 100644 --- a/trace_replay_tester.py +++ b/trace_replay_tester.py @@ -1824,7 +1824,8 @@ async def detect_model(self) -> str: return self.model - def _build_request_params(self, messages: List[dict], max_tokens: Optional[int], stream: bool) -> dict: + def _build_request_params(self, messages: List[dict], max_tokens: Optional[int], stream: bool, + trace_headers: Optional[Dict[str, str]] = None) -> dict: """Build request parameters including model-specific settings.""" params = { "model": self.model, @@ -1835,6 +1836,8 @@ def _build_request_params(self, messages: List[dict], max_tokens: Optional[int], params["max_tokens"] = max_tokens if stream: params["stream_options"] = {"include_usage": True} + if trace_headers: + params["extra_headers"] = trace_headers # Add generation parameters (vLLM supports all of these directly) extra_body = {} @@ -1862,7 +1865,8 @@ def _build_request_params(self, messages: List[dict], max_tokens: Optional[int], async def send_request(self, messages: List[dict], max_tokens: Optional[int], stream: bool = True, on_first_token: Optional[callable] = None, on_chunk: Optional[callable] = None, - tokenizer=None) -> dict: + tokenizer=None, + trace_headers: Optional[Dict[str, str]] = None) -> dict: """ Send request and return metrics. @@ -1902,7 +1906,7 @@ async def send_request(self, messages: List[dict], max_tokens: Optional[int], st prompt_token_ids = None try: - params = self._build_request_params(messages, max_tokens, stream) + params = self._build_request_params(messages, max_tokens, stream, trace_headers=trace_headers) if stream: response = await self.client.chat.completions.create(**params) @@ -2817,7 +2821,12 @@ def on_chunk(chunk_time, chunk_tokens): stream=stream, on_first_token=on_first_token, on_chunk=on_chunk, - tokenizer=self.generator.tokenizer + tokenizer=self.generator.tokenizer, + trace_headers={ + "x-mori-session-id": user.user_id, + "x-mori-trace-id": user.trace_id, + "x-mori-request-idx": str(user.current_idx + 1), + } ) # Decrement decode counter if first token was received From ece52b5a1e5b27a1b6feb1509eb787a849fa0f3a Mon Sep 17 00:00:00 2001 From: AMD-yanfeiwang Date: Thu, 14 May 2026 12:11:31 +0000 Subject: [PATCH 4/4] Align MORI trace replay headers --- trace_replay_tester.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/trace_replay_tester.py b/trace_replay_tester.py index 21d52df..2b4be4d 100644 --- a/trace_replay_tester.py +++ b/trace_replay_tester.py @@ -2823,9 +2823,9 @@ def on_chunk(chunk_time, chunk_tokens): on_chunk=on_chunk, tokenizer=self.generator.tokenizer, trace_headers={ - "x-mori-session-id": user.user_id, - "x-mori-trace-id": user.trace_id, - "x-mori-request-idx": str(user.current_idx + 1), + "x-mori-session-id": user.trace_id, + "x-mori-request-id": str(user.current_idx + 1), + "x-mori-user-id": user.user_id, } )