File size: 7,873 Bytes
e1acc15 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | """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 after __future__ if present; these bases have no futures, but support them.
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:
# Replace only the call prefix, leaving arguments and their evaluation intact.
start = offsets[node.lineno - 1] + node.col_offset
end = offsets[node.func.end_lineno - 1] + node.func.end_col_offset # type: ignore[operator]
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 # type: ignore[operator]
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':
# Decorate the body through a span, preserving returns, exceptions and original signature.
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
# span implementation below uses current-stream events, inclusive of transfer/wait.
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'))
# Time the whole metadata+model interval: start at begin; close in execution finally.
# Verify expected anchors so a foreign/new source cannot silently lose instrumentation.
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()
|