| 1 | #!/usr/bin/env python3 |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | """Which source file produced each GPU kernel symbol in this build. |
| 5 | |
| 6 | Runs inside the csrc-build image stage, after the CMake build, and reads the |
| 7 | objects it left behind. For every compiled object it records the source file, |
| 8 | the CMake target, the headers the compiler reported it depended on (from |
| 9 | ninja's dependency log), and, for device objects, the kernel entry symbols |
| 10 | `cuobjdump` finds in the embedded fatbin. Those mangled names are exactly |
| 11 | what a CUPTI kernel recorder sees at run time, so CI can join "this step |
| 12 | launched kernel K" with "K came from csrc/foo.cu" and select tests for a |
| 13 | kernel change by evidence instead of by image membership. |
| 14 | |
| 15 | kernel_symbol_map.py --build-root build --source-root . \\ |
| 16 | --out dist/kernel_symbol_map.json.gz |
| 17 | |
| 18 | Never fails the image build. If there is nothing to read (a precompiled-wheel |
| 19 | build, a missing cuobjdump) it writes a map with an empty object list and a |
| 20 | `reason`, which the consumer treats as "no evidence", not as "no kernels". |
| 21 | |
| 22 | Output (gzipped JSON): |
| 23 | |
| 24 | { |
| 25 | "version": 1, |
| 26 | "commit": "...", # from --commit or the environment |
| 27 | "cuda": "13.0", # nvcc release when found |
| 28 | "generated_at": "2026-09-19T...", |
| 29 | "build_dirs": ["build/temp.linux-x86_64-cpython-312"], |
| 30 | "incomplete": false, # true when any object's extraction failed; |
| 31 | # objects is then [] and reason is set |
| 32 | "errors": [], # [{"object", "source", "error"}, ...] |
| 33 | "objects": [ |
| 34 | { |
| 35 | "source": "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu", |
| 36 | "target": "_C", |
| 37 | "object": "CMakeFiles/_C.dir/csrc/.../fused_qknorm_rope_kernel.cu.o", |
| 38 | "device": true, |
| 39 | "symbols": ["_Z21fusedQKNormRopeKernelI...", ...], |
| 40 | "deps": ["csrc/libtorch_stable/fused_qknorm_rope_kernel.cu", |
| 41 | "csrc/cuda_compat.h", ...] |
| 42 | # "error": "..." # present when cuobjdump failed on it |
| 43 | }, |
| 44 | ... |
| 45 | ], |
| 46 | "stats": {...}, |
| 47 | "reason": "..." # only when the map is empty |
| 48 | } |
| 49 | |
| 50 | The consumer's contract is deliberately simple: an empty `objects` list with |
| 51 | a `reason` means "no evidence for any file", fall back to the static rule. |
| 52 | A partial map would need every consumer to know which files the unreadable |
| 53 | objects touched, and one that forgot would drop tests that a kernel change |
| 54 | can reach. So the producer never publishes a partial map: if any object is |
| 55 | still unreadable after a retry, `objects` is emptied, `incomplete` is set |
| 56 | and `errors` names the objects so the failure can be diagnosed. |
| 57 | |
| 58 | Paths are relative to --source-root when the file lives inside it, otherwise |
| 59 | absolute (FetchContent sources under the build directory, CUDA headers). |
| 60 | """ |
| 61 | |
| 62 | from __future__ import annotations |
| 63 | |
| 64 | import argparse |
| 65 | import concurrent.futures |
| 66 | import datetime as dt |
| 67 | import gzip |
| 68 | import json |
| 69 | import os |
| 70 | import shutil |
| 71 | import subprocess |
| 72 | import sys |
| 73 | import time |
| 74 | from pathlib import Path |
| 75 | |
| 76 | import regex as re |
| 77 | |
| 78 | DEVICE_SUFFIXES = (".cu.o", ".hip.o") |
| 79 | SOURCE_SUFFIXES = (".cu", ".hip", ".cpp", ".cc", ".cxx", ".c") |
| 80 | DEPS_HEADER = re.compile(r"^(\S+): #deps (\d+), deps mtime \d+ \((VALID|STALE)\)$") |
| 81 | # FetchContent sub-builds nest their own CMakeFiles/ under _deps/<name>-build/. |
| 82 | OBJ_DIR = re.compile(r"(?:^|/)CMakeFiles/([^/]+)\.dir/(.+)\.o$") |
| 83 | |
| 84 | |
| 85 | def log(msg: str) -> None: |
| 86 | print(f"kernel_symbol_map: {msg}", flush=True) |
| 87 | |
| 88 | |
| 89 | def find_tool(name: str, *candidates: str) -> str | None: |
| 90 | for c in candidates: |
| 91 | if c and Path(c).is_file(): |
| 92 | return c |
| 93 | return shutil.which(name) |
| 94 | |
| 95 | |
| 96 | def cuda_release(nvcc: str | None) -> str: |
| 97 | if not nvcc: |
| 98 | return os.environ.get("CUDA_VERSION", "") |
| 99 | try: |
| 100 | out = subprocess.run( |
| 101 | [nvcc, "--version"], capture_output=True, text=True, timeout=30 |
| 102 | ).stdout |
| 103 | m = re.search(r"release (\d+\.\d+)", out) |
| 104 | return m.group(1) if m else os.environ.get("CUDA_VERSION", "") |
| 105 | except (OSError, subprocess.SubprocessError): |
| 106 | return os.environ.get("CUDA_VERSION", "") |
| 107 | |
| 108 | |
| 109 | def find_build_dirs(build_root: Path) -> list[Path]: |
| 110 | """Directories holding a top-level build.ninja (one per CMake build tree).""" |
| 111 | return sorted( |
| 112 | { |
| 113 | p.parent |
| 114 | for p in build_root.rglob("build.ninja") |
| 115 | if (p.parent / ".ninja_deps").exists() or (p.parent / "CMakeFiles").is_dir() |
| 116 | } |
| 117 | ) |
| 118 | |
| 119 | |
| 120 | def ninja_deps(ninja: str | None, build_dir: Path) -> dict[str, list[str]]: |
| 121 | """Object (relative to build_dir) -> paths the compiler reported it read.""" |
| 122 | if not ninja or not (build_dir / ".ninja_deps").exists(): |
| 123 | return {} |
| 124 | try: |
| 125 | out = subprocess.run( |
| 126 | [ninja, "-C", str(build_dir), "-t", "deps"], |
| 127 | capture_output=True, |
| 128 | text=True, |
| 129 | timeout=600, |
| 130 | ).stdout |
| 131 | except (OSError, subprocess.SubprocessError) as e: |
| 132 | log(f"ninja -t deps failed in {build_dir}: {e}") |
| 133 | return {} |
| 134 | deps: dict[str, list[str]] = {} |
| 135 | cur: list[str] | None = None |
| 136 | for line in out.splitlines(): |
| 137 | m = DEPS_HEADER.match(line) |
| 138 | if m: |
| 139 | cur = deps.setdefault(m.group(1), []) |
| 140 | continue |
| 141 | if cur is not None and line.startswith(" "): |
| 142 | cur.append(line.strip()) |
| 143 | elif not line.strip(): |
| 144 | cur = None |
| 145 | return deps |
| 146 | |
| 147 | |
| 148 | def parse_symbols(text: str) -> list[str]: |
| 149 | """Kernel entry points from `cuobjdump -symbols` output. |
| 150 | |
| 151 | Lines look like `STT_FUNC STB_GLOBAL STO_ENTRY _Z21fusedQKNormRope...`. |
| 152 | Entry points are what the driver launches, so they are what a recorder |
| 153 | sees. If a toolchain prints no STO_ENTRY marker, fall back to every global |
| 154 | function symbol; over-attribution is the safe direction. |
| 155 | """ |
| 156 | entries: set[str] = set() |
| 157 | globals_: set[str] = set() |
| 158 | for line in text.splitlines(): |
| 159 | parts = line.split() |
| 160 | if len(parts) < 4 or parts[0] != "STT_FUNC": |
| 161 | continue |
| 162 | if parts[2] == "STO_ENTRY": |
| 163 | entries.add(parts[3]) |
| 164 | elif parts[1] == "STB_GLOBAL": |
| 165 | globals_.add(parts[3]) |
| 166 | return sorted(entries or globals_) |
| 167 | |
| 168 | |
| 169 | def device_symbols(cuobjdump: str, obj: Path) -> tuple[list[str], str | None]: |
| 170 | """(symbols, error). An error means the object is unknown, not kernel-free.""" |
| 171 | syms, err = _cuobjdump_symbols(cuobjdump, obj) |
| 172 | if err: # transient failures happen on large fatbins; one retry is cheap |
| 173 | syms, err = _cuobjdump_symbols(cuobjdump, obj) |
| 174 | return syms, err |
| 175 | |
| 176 | |
| 177 | def _cuobjdump_symbols(cuobjdump: str, obj: Path) -> tuple[list[str], str | None]: |
| 178 | try: |
| 179 | r = subprocess.run( |
| 180 | [cuobjdump, "-symbols", str(obj)], |
| 181 | capture_output=True, |
| 182 | text=True, |
| 183 | timeout=900, |
| 184 | ) |
| 185 | except subprocess.TimeoutExpired: |
| 186 | return [], "timeout" |
| 187 | except OSError as e: |
| 188 | return [], str(e) |
| 189 | if r.returncode != 0: |
| 190 | msg = (r.stderr or r.stdout).strip().splitlines() |
| 191 | return [], (msg[-1] if msg else f"exit {r.returncode}") |
| 192 | return parse_symbols(r.stdout), None |
| 193 | |
| 194 | |
| 195 | def relpath(p: str | Path, source_root: Path) -> str: |
| 196 | path = Path(p) |
| 197 | try: |
| 198 | return path.resolve().relative_to(source_root).as_posix() |
| 199 | except ValueError: |
| 200 | return path.as_posix() |
| 201 | |
| 202 | |
| 203 | def source_of( |
| 204 | obj_rel: str, deps: list[str], build_dir: Path, source_root: Path |
| 205 | ) -> tuple[str, str]: |
| 206 | """(source path, cmake target) for an object. |
| 207 | |
| 208 | The dependency log lists the source first, which is the reliable answer. |
| 209 | Failing that, CMake's object layout `CMakeFiles/<target>.dir/<source>.o` |
| 210 | names it, with `__/` segments standing for `..`. |
| 211 | """ |
| 212 | target = "" |
| 213 | m = OBJ_DIR.search(obj_rel) # search: sub-builds nest CMakeFiles/ |
| 214 | if m: |
| 215 | target = m.group(1) |
| 216 | for d in deps: |
| 217 | if d.endswith(SOURCE_SUFFIXES): |
| 218 | return relpath( |
| 219 | d if os.path.isabs(d) else build_dir / d, source_root |
| 220 | ), target |
| 221 | if m: |
| 222 | guess = m.group(2).replace("__/", "../") |
| 223 | if guess.endswith(SOURCE_SUFFIXES): |
| 224 | for base in (source_root, build_dir): |
| 225 | cand = base / guess |
| 226 | if cand.exists(): |
| 227 | return relpath(cand, source_root), target |
| 228 | return guess, target |
| 229 | return obj_rel, target |
| 230 | |
| 231 | |
| 232 | def main() -> int: |
| 233 | ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) |
| 234 | ap.add_argument("--build-root", type=Path, default=Path("build")) |
| 235 | ap.add_argument("--source-root", type=Path, default=Path(".")) |
| 236 | ap.add_argument("--out", type=Path, required=True) |
| 237 | ap.add_argument("--jobs", type=int, default=min(8, os.cpu_count() or 2)) |
| 238 | ap.add_argument("--cuobjdump", default=None) |
| 239 | ap.add_argument("--commit", default="", help="commit the build is of") |
| 240 | a = ap.parse_args() |
| 241 | t0 = time.time() |
| 242 | source_root = a.source_root.resolve() |
| 243 | |
| 244 | cuda_home = os.environ.get("CUDA_HOME", "/usr/local/cuda") |
| 245 | cuobjdump = a.cuobjdump or find_tool( |
| 246 | "cuobjdump", f"{cuda_home}/bin/cuobjdump", "/usr/local/cuda/bin/cuobjdump" |
| 247 | ) |
| 248 | nvcc = find_tool("nvcc", f"{cuda_home}/bin/nvcc", "/usr/local/cuda/bin/nvcc") |
| 249 | ninja = shutil.which("ninja") |
| 250 | |
| 251 | result: dict = { |
| 252 | "version": 1, |
| 253 | "commit": a.commit |
| 254 | or os.environ.get("BUILDKITE_COMMIT") |
| 255 | or os.environ.get("VLLM_BUILD_COMMIT") |
| 256 | or "", |
| 257 | "cuda": cuda_release(nvcc), |
| 258 | "generated_at": dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds"), |
| 259 | "build_dirs": [], |
| 260 | "incomplete": False, |
| 261 | "errors": [], |
| 262 | "objects": [], |
| 263 | "stats": {}, |
| 264 | } |
| 265 | |
| 266 | build_dirs = find_build_dirs(a.build_root) if a.build_root.exists() else [] |
| 267 | result["build_dirs"] = [relpath(b, source_root) for b in build_dirs] |
| 268 | if not build_dirs: |
| 269 | result["reason"] = ( |
| 270 | f"no CMake build tree under {a.build_root} (precompiled wheel build?)" |
| 271 | ) |
| 272 | elif not cuobjdump: |
| 273 | result["reason"] = "cuobjdump not found" |
| 274 | log( |
| 275 | f"build dirs: {result['build_dirs'] or 'none'}; " |
| 276 | f"cuobjdump={cuobjdump}; ninja={ninja}" |
| 277 | ) |
| 278 | |
| 279 | jobs: list[tuple[Path, str, Path]] = [] # (build_dir, obj_rel, obj_path) |
| 280 | all_deps: dict[Path, dict[str, list[str]]] = {} |
| 281 | if "reason" not in result: |
| 282 | for bd in build_dirs: |
| 283 | all_deps[bd] = ninja_deps(ninja, bd) |
| 284 | for obj in bd.rglob("*.o"): |
| 285 | rel = obj.relative_to(bd).as_posix() |
| 286 | # CMake's compiler-identification probes are objects too. |
| 287 | if "CompilerId" in rel: |
| 288 | continue |
| 289 | if "/CMakeFiles/" in f"/{rel}" or rel.startswith("CMakeFiles/"): |
| 290 | jobs.append((bd, rel, obj)) |
| 291 | n_deps = sum(len(d) for d in all_deps.values()) |
| 292 | log(f"{len(jobs)} objects, {n_deps} dependency records") |
| 293 | |
| 294 | def work(job): |
| 295 | bd, rel, obj = job |
| 296 | deps = all_deps.get(bd, {}).get(rel, []) |
| 297 | source, target = source_of(rel, deps, bd, source_root) |
| 298 | device = rel.endswith(DEVICE_SUFFIXES) |
| 299 | syms, err = device_symbols(cuobjdump, obj) if device else ([], None) |
| 300 | rel_deps = sorted( |
| 301 | {relpath(d if os.path.isabs(d) else bd / d, source_root) for d in deps} |
| 302 | ) |
| 303 | entry = { |
| 304 | "source": source, |
| 305 | "target": target, |
| 306 | "object": rel, |
| 307 | "device": device, |
| 308 | "symbols": syms, |
| 309 | "deps": rel_deps, |
| 310 | } |
| 311 | if err: |
| 312 | entry["error"] = err |
| 313 | return entry |
| 314 | |
| 315 | with concurrent.futures.ThreadPoolExecutor(max_workers=max(1, a.jobs)) as ex: |
| 316 | result["objects"] = sorted( |
| 317 | ex.map(work, jobs), key=lambda e: (e["source"], e["target"]) |
| 318 | ) |
| 319 | |
| 320 | # An object cuobjdump could not read is unknown, not kernel-free. Say so |
| 321 | # at the top so a consumer never has to scan for it, and name the objects |
| 322 | # so it can widen the "unknown" set to every file they touch. |
| 323 | result["errors"] = [ |
| 324 | {"object": e["object"], "source": e["source"], "error": e["error"]} |
| 325 | for e in result["objects"] |
| 326 | if "error" in e |
| 327 | ] |
| 328 | result["incomplete"] = bool(result["errors"]) |
| 329 | for e in result["errors"][:5]: |
| 330 | log(f"cuobjdump failed on {e['object']}: {e['error']}") |
| 331 | if result["incomplete"]: |
| 332 | # Publish nothing rather than something partial: see the docstring. |
| 333 | result["objects"] = [] |
| 334 | result["reason"] = ( |
| 335 | f"{len(result['errors'])} objects unreadable by cuobjdump, e.g. " |
| 336 | f"{result['errors'][0]['object']}: {result['errors'][0]['error']}" |
| 337 | ) |
| 338 | |
| 339 | result["stats"] = { |
| 340 | "objects": len(jobs), |
| 341 | "device_objects": sum(1 for e in result["objects"] if e["device"]), |
| 342 | "objects_with_symbols": sum(1 for e in result["objects"] if e["symbols"]), |
| 343 | "symbols": sum(len(e["symbols"]) for e in result["objects"]), |
| 344 | "cuobjdump_errors": len(result["errors"]), |
| 345 | "seconds": round(time.time() - t0, 1), |
| 346 | } |
| 347 | |
| 348 | a.out.parent.mkdir(parents=True, exist_ok=True) |
| 349 | with gzip.open(a.out, "wt", encoding="utf-8") as f: |
| 350 | json.dump(result, f, separators=(",", ":")) |
| 351 | note = f"; reason: {result['reason']}" if "reason" in result else "" |
| 352 | log(f"wrote {a.out} ({a.out.stat().st_size // 1024} KiB): {result['stats']}{note}") |
| 353 | |
| 354 | # One line a human can check in the build log. |
| 355 | probe = [ |
| 356 | e |
| 357 | for e in result["objects"] |
| 358 | if any("fusedQKNormRopeKernel" in s for s in e["symbols"]) |
| 359 | ] |
| 360 | if probe: |
| 361 | log(f"probe: fusedQKNormRopeKernel -> {sorted({e['source'] for e in probe})}") |
| 362 | return 0 |
| 363 | |
| 364 | |
| 365 | if __name__ == "__main__": |
| 366 | try: |
| 367 | sys.exit(main()) |
| 368 | except Exception as e: # never fail the image build |
| 369 | log(f"unexpected error, continuing without a map: {e!r}") |
| 370 | sys.exit(0) |