Freeze Qwen30 trace profile support
This commit is contained in:
@@ -0,0 +1,215 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Freeze the exact request cohort and its operator-profile support.
|
||||
|
||||
This calls AITuner's production trace loader, including its input-length
|
||||
filter, uniform max-request downsampling, output override, and sampling-u
|
||||
threshold semantics. Prefix reuse below is a no-eviction upper bound; the
|
||||
actual cache state remains scheduler/config dependent.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable
|
||||
|
||||
from aituner.spec import load_study_spec
|
||||
from aituner.trace import load_trace_requests, select_requests_for_threshold
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--study", type=Path, required=True)
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument("--cohort-output", type=Path, required=True)
|
||||
parser.add_argument(
|
||||
"--thresholds",
|
||||
type=float,
|
||||
nargs="+",
|
||||
default=[0.125, 0.25, 0.5, 1.0],
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def sha256_file(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def nearest_rank(values: Iterable[float], percentile: float) -> float:
|
||||
ordered = sorted(values)
|
||||
if not ordered:
|
||||
return 0.0
|
||||
index = min(
|
||||
len(ordered) - 1,
|
||||
max(0, math.ceil(percentile / 100.0 * len(ordered)) - 1),
|
||||
)
|
||||
return float(ordered[index])
|
||||
|
||||
|
||||
def distribution(values: list[float]) -> dict[str, float]:
|
||||
return {
|
||||
"min": min(values, default=0.0),
|
||||
"p10": nearest_rank(values, 10),
|
||||
"p25": nearest_rank(values, 25),
|
||||
"p50": nearest_rank(values, 50),
|
||||
"p75": nearest_rank(values, 75),
|
||||
"p90": nearest_rank(values, 90),
|
||||
"p95": nearest_rank(values, 95),
|
||||
"p99": nearest_rank(values, 99),
|
||||
"max": max(values, default=0.0),
|
||||
}
|
||||
|
||||
|
||||
def prefix_upper_bound(requests: list[Any], block_size: int) -> dict[str, Any]:
|
||||
seen: set[Any] = set()
|
||||
total_blocks = 0
|
||||
reusable_blocks = 0
|
||||
leading_reusable_blocks: list[float] = []
|
||||
reusable_tokens: list[float] = []
|
||||
rows_with_hashes = 0
|
||||
for request in requests:
|
||||
hashes = request.metadata.get("hash_ids")
|
||||
if not isinstance(hashes, list):
|
||||
continue
|
||||
rows_with_hashes += 1
|
||||
total_blocks += len(hashes)
|
||||
reusable = sum(hash_id in seen for hash_id in hashes)
|
||||
reusable_blocks += reusable
|
||||
leading = 0
|
||||
for hash_id in hashes:
|
||||
if hash_id not in seen:
|
||||
break
|
||||
leading += 1
|
||||
leading_reusable_blocks.append(float(leading))
|
||||
reusable_tokens.append(
|
||||
float(min(request.prompt_tokens_hint or 0, leading * block_size))
|
||||
)
|
||||
seen.update(hashes)
|
||||
return {
|
||||
"semantics": (
|
||||
"arrival-ordered, infinite-capacity/no-eviction upper bound; "
|
||||
"not an observed KV-cache hit rate"
|
||||
),
|
||||
"rows_with_hash_ids": rows_with_hashes,
|
||||
"total_blocks": total_blocks,
|
||||
"unique_blocks": len(seen),
|
||||
"any_position_reusable_block_ratio": (
|
||||
reusable_blocks / total_blocks if total_blocks else 0.0
|
||||
),
|
||||
"leading_reusable_blocks_per_request": distribution(
|
||||
leading_reusable_blocks
|
||||
),
|
||||
"leading_reusable_tokens_per_request": distribution(reusable_tokens),
|
||||
}
|
||||
|
||||
|
||||
def summarize(requests: list[Any], block_size: int) -> dict[str, Any]:
|
||||
input_lengths = [float(request.prompt_tokens_hint or 0) for request in requests]
|
||||
output_lengths = [
|
||||
float(request.completion_tokens_hint or 0) for request in requests
|
||||
]
|
||||
hash_counts = [
|
||||
float(len(request.metadata.get("hash_ids") or [])) for request in requests
|
||||
]
|
||||
arrivals = [request.arrival_s for request in requests]
|
||||
interarrivals = [
|
||||
max(0.0, arrivals[index] - arrivals[index - 1])
|
||||
for index in range(1, len(arrivals))
|
||||
]
|
||||
return {
|
||||
"request_count": len(requests),
|
||||
"input_tokens": distribution(input_lengths),
|
||||
"output_tokens": distribution(output_lengths),
|
||||
"hash_blocks_per_request": distribution(hash_counts),
|
||||
"interarrival_s": distribution(interarrivals),
|
||||
"sampling_u": distribution([request.sampling_u for request in requests]),
|
||||
"multi_turn_fraction": (
|
||||
sum(
|
||||
isinstance(request.metadata.get("turn"), (int, float))
|
||||
and request.metadata["turn"] > 1
|
||||
for request in requests
|
||||
)
|
||||
/ len(requests)
|
||||
if requests
|
||||
else 0.0
|
||||
),
|
||||
"prefix_reuse_upper_bound": prefix_upper_bound(requests, block_size),
|
||||
}
|
||||
|
||||
|
||||
def cohort_row(request: Any) -> dict[str, Any]:
|
||||
return {
|
||||
"row_id": request.row_id,
|
||||
"arrival_s": request.arrival_s,
|
||||
"sampling_u": request.sampling_u,
|
||||
"input_tokens": request.prompt_tokens_hint,
|
||||
"output_tokens": request.completion_tokens_hint,
|
||||
"hash_ids": request.metadata.get("hash_ids"),
|
||||
"turn": request.metadata.get("turn"),
|
||||
"parent_chat_id": request.metadata.get("parent_chat_id"),
|
||||
"type": request.metadata.get("type"),
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
study = load_study_spec(args.study)
|
||||
window, cohort = load_trace_requests(study, study_spec_path=args.study)
|
||||
block_size = int(window.source_payload.get("block_size") or 1)
|
||||
|
||||
args.cohort_output.parent.mkdir(parents=True, exist_ok=True)
|
||||
with args.cohort_output.open("w", encoding="utf-8") as handle:
|
||||
for request in cohort:
|
||||
handle.write(json.dumps(cohort_row(request), sort_keys=True) + "\n")
|
||||
|
||||
payload = {
|
||||
"schema_version": "qwen30_trace_profile_support.v1",
|
||||
"study": str(args.study.resolve()),
|
||||
"study_sha256": sha256_file(args.study),
|
||||
"trace": str(window.trace_path),
|
||||
"trace_sha256": sha256_file(window.trace_path),
|
||||
"window_id": window.window_id,
|
||||
"block_size": block_size,
|
||||
"loader_contract": {
|
||||
"input_length_filter": {
|
||||
"min": study.trace.input_length_filter.min_input_tokens,
|
||||
"max": study.trace.input_length_filter.max_input_tokens,
|
||||
}
|
||||
if study.trace.input_length_filter is not None
|
||||
else None,
|
||||
"completion_tokens_override": study.trace.completion_tokens_override,
|
||||
"max_requests_per_probe": study.trace.max_requests_per_probe,
|
||||
"replay_time_scale": study.trace.replay_time_scale,
|
||||
"ordering": "arrival_s",
|
||||
"downsampling": "AITuner _downsample_requests before sampling_u threshold",
|
||||
},
|
||||
"full_downsampled_cohort": summarize(cohort, block_size),
|
||||
"threshold_cohorts": {
|
||||
str(threshold): summarize(
|
||||
select_requests_for_threshold(cohort, threshold=threshold),
|
||||
block_size,
|
||||
)
|
||||
for threshold in args.thresholds
|
||||
},
|
||||
"limits": [
|
||||
"Trace length/hash support does not determine dynamic decode or mixed batch shapes.",
|
||||
"Those shapes depend jointly on arrival history, SLO pressure, TP execution time, MNS, chunking, and KV eviction.",
|
||||
"MoE expert routing is not present in the trace and must be measured from model execution.",
|
||||
],
|
||||
}
|
||||
payload["cohort_output"] = str(args.cohort_output.resolve())
|
||||
payload["cohort_sha256"] = sha256_file(args.cohort_output)
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
|
||||
print(json.dumps(payload, indent=2, sort_keys=True))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user