| """Build a hash-pinned diagnostic overlay. Default OFF emits byte-identical bases.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import ast |
| import hashlib |
| import json |
| from pathlib import Path |
| from typing import cast |
|
|
| HERE = Path(__file__).resolve().parent |
| REPO = HERE.parents[1] |
| FILES = { |
| 'vllm/v1/worker/gpu/model_runner.py': 'k1', |
| 'vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py': 'k1', |
| 'vllm/model_executor/layers/sparse_attn_indexer.py': 'image', |
| 'vllm/v1/attention/backends/mla/indexer.py': 'image', |
| 'vllm/distributed/parallel_state.py': 'image', |
| } |
| HELPER = 'vllm/v1/worker/gpu/dsv41_prefill_trace.py' |
|
|
|
|
| def sha(data: bytes) -> str: |
| return hashlib.sha256(data).hexdigest() |
|
|
|
|
| def instrument(text: str, path: str) -> str: |
| tree = ast.parse(text) |
| lines = text.splitlines(keepends=True) |
| offsets = [0] |
| for line in lines: |
| offsets.append(offsets[-1] + len(line)) |
| edits: list[tuple[int, int, str]] = [] |
| |
| imports = [n for n in tree.body if isinstance(n, ast.ImportFrom) and n.module == '__future__'] |
| at = offsets[imports[-1].end_lineno or 0] if imports else offsets[tree.body[0].end_lineno or 0] |
| edits.append((at, at, '\nfrom vllm.v1.worker.gpu import dsv41_prefill_trace as _pf\n')) |
| targets = { |
| 'cp_gather_indexer_k_quant_cache': 'indexer.gather', |
| 'fp8_fp4_mqa_logits': 'indexer.mqa', |
| 'top_k_per_row_prefill': 'indexer.topk', |
| '_select_candidate_blocks': 'indexer.candidate_select', |
| '_apply_candidate_mask': 'indexer.candidate_mask', |
| 'exchange_topk': 'collective.indexer_exchange', |
| 'compute_global_topk_indices_and_lens': 'sparse.metadata_scan', |
| 'sparse_attention': 'sparse.dispatch.triton', |
| 'flashinfer_trtllm_batch_decode_sparse_mla_dsv4': 'sparse.dispatch.flashinfer', |
| 'prepare_attn': 'metadata.prepare', |
| 'build_prefill_chunk_metadata': 'metadata.indexer_chunk', |
| } |
| hits: dict[str, int] = {} |
| for node in ast.walk(tree): |
| if isinstance(node, ast.Call): |
| name = node.func.attr if isinstance(node.func, ast.Attribute) else node.func.id if isinstance(node.func, ast.Name) else '' |
| if name in targets: |
| |
| start = offsets[node.lineno - 1] + node.col_offset |
| end = offsets[node.func.end_lineno - 1] + node.func.end_col_offset |
| original_fn = text[start:end] |
| opening = text.index('(', end) |
| has_args = bool(node.args or node.keywords) |
| edits.append((start, opening + 1, |
| f'_pf.call({targets[name]!r}, _pf.meta(locals()), {original_fn}' + (', ' if has_args else ''))) |
| hits[name] = hits.get(name, 0) + 1 |
| if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): |
| if node.name == 'execute_model' and path.endswith('model_runner.py'): |
| pos = offsets[node.lineno - 1] + node.col_offset |
| edits.append((pos, pos, '@_pf.execution\n' + ' ' * node.col_offset)) |
| if node.name == '_split_indexer_prefill_chunks' and path.endswith('/mla/indexer.py'): |
| returns = [n for n in ast.walk(node) if isinstance(n, ast.Return)] |
| if len(returns) != 1: |
| raise ValueError('PF1 unexpected planner returns') |
| ret = returns[0] |
| start = offsets[ret.lineno - 1] + ret.col_offset |
| end = offsets[ret.end_lineno - 1] + ret.end_col_offset |
| edits.append((start, end, 'return _pf.plan(chunks, compressed_seq_lens_cpu, prefill_query_lens_cpu, workspace_size, max_logits_bytes, request_offset)')) |
| if path.endswith('parallel_state.py') and node.name in ('all_reduce', 'all_gather', 'reduce_scatter', 'all_gatherv', 'reduce_scatterv') and node.args.args and node.args.args[0].arg == 'self': |
| |
| body = node.body[1:] if isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Constant) else node.body |
| start = offsets[body[0].lineno - 1] |
| end = offsets[node.end_lineno or node.lineno] |
| original = text[start:end] |
| indent = ' ' * body[0].col_offset |
| |
| edits.append((start, end, indent + f'with _pf.span("collective.{node.name}", _pf.meta(locals())):\n' + ''.join(' ' + ln if ln.strip() else ln for ln in original.splitlines(keepends=True)))) |
| if path.endswith('model_runner.py'): |
| anchor = ' attn_metadata = None\n' |
| if text.count(anchor) != 1: |
| raise ValueError('PF1 unexpected runner anchor') |
| pos = text.index(anchor) |
| edits.append((pos, pos, ' _pf.begin(input_batch, dummy_run, batch_desc.cg_mode, self.max_model_len)\n')) |
| |
| |
| required = {'model_runner.py': ('prepare_attn',), 'flashinfer_sparse.py': ('flashinfer_trtllm_batch_decode_sparse_mla_dsv4',), |
| 'sparse_attn_indexer.py': ('cp_gather_indexer_k_quant_cache', 'fp8_fp4_mqa_logits', 'top_k_per_row_prefill'), |
| 'indexer.py': ('build_prefill_chunk_metadata',)} |
| for key in required.get(Path(path).name, ()): |
| if key not in hits: |
| raise ValueError(f'PF1 missing instrumentation anchor {path}:{key}') |
| edits.sort(reverse=True) |
| last_start = len(text) + 1 |
| for start, end, replacement in edits: |
| if end > last_start: |
| raise ValueError(f'PF1 overlapping instrumentation edits in {path}') |
| text = text[:start] + replacement + text[end:] |
| last_start = start |
| compile(text, path, 'exec') |
| return text |
|
|
|
|
| def build(output: Path, image: Path, enabled: bool = False) -> dict[str, str]: |
| pins = cast(dict[str, str], json.loads((HERE / 'BASES.json').read_text())) |
| if output.resolve().is_relative_to(image.resolve()) or output.resolve().is_relative_to((REPO / 'offload/kvpack/overlay').resolve()): |
| raise ValueError('PF1 output must not modify reference/production overlay') |
| sources: dict[str, bytes] = {} |
| for path, source in FILES.items(): |
| base = (REPO / 'offload/kvpack/overlay' if source == 'k1' else image) / path |
| data = base.read_bytes() |
| if sha(data) != pins[path]: |
| raise ValueError(f'PF1 trace base SHA mismatch: {base}') |
| sources[path] = instrument(data.decode(), path).encode() if enabled else data |
| if enabled: |
| sources[HELPER] = (HERE / 'runtime.py').read_bytes() |
| hashes: dict[str, str] = {} |
| for path, data in sources.items(): |
| target = output / path |
| target.parent.mkdir(parents=True, exist_ok=True) |
| if target.exists() and target.read_bytes() != data: |
| raise ValueError(f'PF1 refuses existing different output: {target}') |
| target.write_bytes(data) |
| hashes[path] = sha(data) |
| return hashes |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser(description=__doc__) |
| ap.add_argument('--output-root', type=Path, required=True) |
| ap.add_argument('--source-root', type=Path, default=Path('/opt/dsv41/data/src-image')) |
| ap.add_argument('--enabled', action='store_true') |
| args = ap.parse_args() |
| print(json.dumps(build(args.output_root, args.source_root, args.enabled), indent=2)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
