| """Window-only numeric telemetry. No tensor contents, device reads or synchronization.""" |
| from __future__ import annotations |
|
|
| import functools |
| import hashlib |
| import json |
| import os |
| import sys |
| import time |
| from collections import deque |
| from collections.abc import Callable, Iterator, Mapping, Sized |
| from contextlib import contextmanager |
| from typing import ContextManager |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import ParamSpec, Protocol, TypeVar, cast |
|
|
| Scalar = str | int | float | bool | None |
| P = ParamSpec('P') |
| T = TypeVar('T') |
| ENABLED = os.environ.get('VLLM_DSV41_PREFILL_TRACE', '0') == '1' |
|
|
|
|
| class Event(Protocol): |
| def record(self) -> None: ... |
| def query(self) -> bool: ... |
| def elapsed_time(self, end: Event) -> float: ... |
|
|
|
|
| class Backend(Protocol): |
| def event(self) -> Event: ... |
| def rank(self) -> int: ... |
| def capturing(self) -> bool: ... |
|
|
|
|
| class CudaBackend: |
| def event(self) -> Event: |
| import torch |
| factory = cast(Callable[..., Event], torch.cuda.Event) |
| return factory(enable_timing=True) |
|
|
| def rank(self) -> int: |
| import torch.distributed as dist |
| return dist.get_rank() |
|
|
| def capturing(self) -> bool: |
| import torch |
| return torch.cuda.is_current_stream_capturing() |
|
|
|
|
| @dataclass |
| class Pending: |
| row: dict[str, object] |
| start: Event |
| end: Event |
|
|
|
|
| class Recorder: |
| def __init__(self, path: Path, backend: Backend, limit: int = 8192) -> None: |
| if limit < 1: |
| raise ValueError('trace pending limit must be positive') |
| self.backend = backend |
| self.limit = limit |
| self.pending: deque[Pending] = deque() |
| self.path = path |
| self.context: dict[str, object] | None = None |
| self.sequence = 0 |
| self.whole: ContextManager[None] | None = None |
| self.occurrences: dict[str, int] = {} |
| |
| self.stream = path.open('x', buffering=1) |
|
|
| def write(self, row: dict[str, object]) -> None: |
| self.stream.write(json.dumps(row, sort_keys=True, allow_nan=False) + '\n') |
|
|
| def drain(self) -> None: |
| |
| remaining: deque[Pending] = deque() |
| while self.pending: |
| item = self.pending.popleft() |
| if item.end.query(): |
| item.row['cuda_ms'] = item.start.elapsed_time(item.end) |
| self.write(item.row) |
| else: |
| remaining.append(item) |
| self.pending = remaining |
|
|
| def call(self, label: str, metadata: dict[str, Scalar], fn: Callable[P, T], |
| *args: P.args, **kwargs: P.kwargs) -> T: |
| if self.context is None: |
| return fn(*args, **kwargs) |
| if self.backend.capturing(): |
| raise RuntimeError('PF1 trace requires eager prefill; graph capture is unsupported') |
| if len(self.pending) >= self.limit: |
| self.drain() |
| if len(self.pending) >= self.limit: |
| raise RuntimeError('PF1 trace pending event bound exceeded; stop window, do not drop samples') |
| start, end = self.backend.event(), self.backend.event() |
| row: dict[str, object] = {**self.context, **metadata, 'phase': label, 'event': self.sequence, |
| 'host_start_ns': time.monotonic_ns()} |
| self.sequence += 1 |
| start.record() |
| begin = time.perf_counter_ns() |
| try: |
| return fn(*args, **kwargs) |
| except BaseException: |
| row['failed'] = True |
| raise |
| finally: |
| row['wall_ms'] = (time.perf_counter_ns() - begin) / 1e6 |
| end.record() |
| self.pending.append(Pending(row, start, end)) |
|
|
|
|
| _recorder: Recorder | None = None |
|
|
|
|
| def recorder() -> Recorder: |
| global _recorder |
| if _recorder is None: |
| root = Path(os.environ['VLLM_DSV41_PREFILL_TRACE_DIR']) |
| if not root.is_dir(): |
| raise RuntimeError(f'PF1 trace directory must already exist: {root}') |
| backend = CudaBackend() |
| _recorder = Recorder(root / f'prefill-r{backend.rank()}-p{os.getpid()}.jsonl', backend) |
| return _recorder |
|
|
|
|
| def attr(value: object, name: str) -> object: |
| return getattr(value, name, None) |
|
|
|
|
| def meta(scope: Mapping[str, object]) -> dict[str, Scalar]: |
| """Read only scalar Python attributes and tensor SHAPES, never GPU values.""" |
| out: dict[str, Scalar] = {} |
| for key in ('row_start', 'row_end', 'topk_tokens', 'num_rows', 'num_tokens', |
| 'num_prefills', 'num_prefill_tokens', 'num_chunks', 'chunk_idx', |
| 'chunk_start', 'chunk_end', 'total_seq_lens', 'head_dim', 'use_fp4_cache', |
| 'candidate_write', 'candidate_block_size', 'topk_backend', 'layer_name'): |
| value = scope.get(key) |
| if isinstance(value, (int, float, bool, str)): |
| out[key] = value |
| owner = scope.get('self') |
| for key in ('compress_ratio', 'max_prefill_buffer_size', 'unique_name', 'prefix'): |
| value = attr(owner, key) |
| if isinstance(value, (int, float, bool, str)): |
| out[key] = value |
| chunk = scope.get('chunk') |
| for key in ('token_start', 'token_end', 'total_seq_lens', 'local_total_seq_lens', |
| 'max_local_total_seq_lens', 'num_reqs', 'skip_kv_gather'): |
| value = attr(chunk, key) |
| if isinstance(value, (int, bool)): |
| out['indexer_chunk_' + key] = value |
| for key in ('logits', 'k_quant', 'k_scale', 'q_chunk', 'q_slice', 'swa_indices_chunk', |
| 'extra_sparse_indices_chunk', 'topk_indices', 'block_table', |
| 'indexer_block_table', 'input_'): |
| value = scope.get(key) |
| shape = attr(value, 'shape') |
| if shape is not None: |
| |
| dims = tuple(cast(tuple[int, ...], shape)) |
| out[key + '_shape'] = 'x'.join(str(int(dim)) for dim in dims) |
| size = 1 |
| for dim in dims: |
| size *= int(dim) |
| out[key + '_elements'] = size |
| width = cast(Callable[[], int], getattr(value, 'element_size'))() |
| out[key + '_bytes'] = size * width |
| cm = scope.get('common_attn_metadata') |
| table = attr(cm, 'block_table_tensor') |
| shape = attr(table, 'shape') |
| if shape is not None: |
| dims = tuple(cast(tuple[int, ...], shape)) |
| out['metadata_block_table_rows'] = int(dims[0]) |
| out['metadata_block_table_columns'] = int(dims[1]) |
| out['metadata_block_table_elements_bound'] = int(dims[0]) * int(dims[1]) |
| return out |
|
|
|
|
| class HostArray(Protocol): |
| def __getitem__(self, key: int) -> object: ... |
|
|
|
|
| class Batch(Protocol): |
| req_ids: list[str] |
| num_reqs: int |
| has_prefill: bool |
| is_prefilling_np: HostArray |
| query_start_loc_np: HostArray |
| num_computed_prefill_tokens_np: HostArray |
|
|
|
|
| def begin(batch: Batch, dummy: bool, mode: object, cap: int = 0) -> None: |
| if not ENABLED: |
| return |
| rec = recorder() |
| rec.drain() |
| rec.context = None |
| if dummy or not batch.has_prefill: |
| return |
| if str(mode).split('.')[-1] != 'NONE': |
| raise RuntimeError('PF1 trace requires CUDAGraphMode.NONE for prefill') |
| members: list[dict[str, object]] = [] |
| salt = os.environ['VLLM_DSV41_PREFILL_TRACE_RUN'] |
| for i in range(batch.num_reqs): |
| |
| if not bool(batch.is_prefilling_np[i]): |
| continue |
| key = hashlib.sha256((salt + ':' + batch.req_ids[i]).encode()).hexdigest()[:24] |
| prefix = int(cast(int, batch.num_computed_prefill_tokens_np[i])) |
| q = int(cast(int, batch.query_start_loc_np[i + 1])) - int(cast(int, batch.query_start_loc_np[i])) |
| identity = f'{key}:{prefix}:{q}' |
| occurrence = rec.occurrences.get(identity, 0) |
| rec.occurrences[identity] = occurrence + 1 |
| if len(rec.occurrences) > 32768: |
| raise RuntimeError('PF1 trace chunk identity bound exceeded; rotate window') |
| members.append({'request': key, 'chunk': identity + f':{occurrence}', |
| 'new_tokens': q, 'prefix_tokens': prefix, |
| 'batch_row_start': int(cast(int, batch.query_start_loc_np[i])), |
| 'batch_row_end': int(cast(int, batch.query_start_loc_np[i + 1]))}) |
| rec.context = {'rank': rec.backend.rank(), 'run': salt, 'max_model_len': cap, 'members': members} |
| rec.write({**rec.context, 'phase': 'chunk.begin', 'host_start_ns': time.monotonic_ns()}) |
| rec.whole = span('chunk.total', {}) |
| rec.whole.__enter__() |
|
|
|
|
| def call(label: str, metadata: dict[str, Scalar], fn: Callable[P, T], |
| *args: P.args, **kwargs: P.kwargs) -> T: |
| if not ENABLED or _recorder is None: |
| return fn(*args, **kwargs) |
| return _recorder.call(label, metadata, fn, *args, **kwargs) |
|
|
|
|
| def execution(fn: Callable[P, T]) -> Callable[P, T]: |
| @functools.wraps(fn) |
| def wrapped(*args: P.args, **kwargs: P.kwargs) -> T: |
| try: |
| return fn(*args, **kwargs) |
| finally: |
| if ENABLED and _recorder is not None: |
| if _recorder.whole is not None: |
| _recorder.whole.__exit__(*sys.exc_info()) |
| _recorder.whole = None |
| if _recorder.context is not None: |
| _recorder.write({**_recorder.context, 'phase': 'chunk.end', |
| 'host_end_ns': time.monotonic_ns(), |
| 'pending_events': len(_recorder.pending)}) |
| _recorder.context = None |
| _recorder.drain() |
| return wrapped |
|
|
|
|
| def plan(chunks: list[tuple[slice, slice]], seq: HostArray, query: HostArray, |
| workspace: int, logits_bytes: int, offset: int) -> list[tuple[slice, slice]]: |
| if ENABLED and _recorder is not None and _recorder.context is not None: |
| rec = _recorder |
| assert rec.context is not None |
| requested_rows = sum(int(cast(int, seq[j])) for j in range(len(cast(Sized, seq)))) |
| for i, (req, rows) in enumerate(chunks): |
| start, stop = cast(int, req.start) - offset, cast(int, req.stop) - offset |
| key_rows = sum(int(cast(int, seq[j])) for j in range(start, stop)) |
| new_rows = sum(int(cast(int, query[j])) for j in range(start, stop)) |
| rec.write({**rec.context, 'phase': 'indexer.plan', 'plan_chunk': i, |
| 'req_start': start + offset, 'req_stop': stop + offset, |
| 'query_start': rows.start, 'query_stop': rows.stop, |
| 'compressed_key_rows': key_rows, 'query_rows_before_split': new_rows, |
| 'workspace_rows': workspace, 'logits_budget_bytes': logits_bytes, |
| 'workspace_exceeded': key_rows > workspace, |
| 'all_requested_key_rows': requested_rows, |
| 'workspace_group_split': requested_rows > workspace and key_rows <= workspace, |
| 'query_subchunked': rows.start != 0 or rows.stop != new_rows, |
| 'fallback': 'single_request_over_workspace' if key_rows > workspace else 'none'}) |
| return chunks |
|
|
| |
| @contextmanager |
| def span(label: str, metadata: dict[str, Scalar]) -> Iterator[None]: |
| if not ENABLED or _recorder is None or _recorder.context is None: |
| yield |
| return |
| rec = _recorder |
| assert rec.context is not None |
| if rec.backend.capturing(): |
| raise RuntimeError('PF1 trace cannot insert events into a captured graph') |
| if len(rec.pending) >= rec.limit: |
| rec.drain() |
| if len(rec.pending) >= rec.limit: |
| raise RuntimeError('PF1 trace pending event bound exceeded') |
| start, end = rec.backend.event(), rec.backend.event() |
| row: dict[str, object] = {**rec.context, **metadata, 'phase': label, 'event': rec.sequence, |
| 'host_start_ns': time.monotonic_ns()} |
| rec.sequence += 1 |
| start.record() |
| begin_ns = time.perf_counter_ns() |
| try: |
| yield |
| except BaseException: |
| row['failed'] = True |
| raise |
| finally: |
| row['wall_ms'] = (time.perf_counter_ns() - begin_ns) / 1e6 |
| end.record() |
| rec.pending.append(Pending(row, start, end)) |
|
|
