[CI][ROCm] Record native kernel coverage for AMD jobs

buildkite/ci_selector/README.md

@@ -10,7 +10,7 @@ It picks the jobs for every vLLM pull request build. **New here, or wondering wh
 
 **The coverage record** is a table of what each step actually ran on real CI builds, one row per step, produced by an instrumented build. The recorder that produces it is `recorders/fnrec/`, loaded into every Python process of a step when the build has `VLLM_CI_FNREC=1`; `scripts/` turns a build's recordings into the table.
 
-**The kernel record** is the same idea for `csrc/`, where no Python frame exists: a table of the GPU kernels each step launched (CUPTI, recorded on the nightly and daily runs) joined to a map of which csrc file each kernel was compiled from (read off the image build's objects). Both are produced by `recorders/kernrec/`.
+**The kernel record** is the same idea for `csrc/`, where no Python frame exists: a table of the GPU kernels each step launched (CUPTI on NVIDIA or [ROCProfiler SDK on AMD](recorders/kernrec/ROCM.md)) joined to a map of which csrc file each kernel was compiled from (read off the image build's objects). `recorders/kernrec/` collects the tables; vLLM's image builds produce the maps.
 
 **The library record** answers for a dependency bump, where the changed file is a version pin: the same recorder also writes, per step, which functions of a few watched libraries (flashinfer, DeepGEMM, flash-attn, Triton; `coverage/libraries.py`) code outside them called at runtime. A diff that only moves such a pin runs the steps that called the library instead of every step on its image. It only drops, only steps selected for such pins alone, and a row recorded before the library was watched keeps its step. Torch is not watched: a torch bump runs everything.
 
@@ -24,7 +24,7 @@ Neither is a stage of the other. `decide.py` reads all of them, per changed file
 
 A changed `.cu` votes with the kernels its diff touched, not the whole file: `coverage/changed_kernels.py` reads both sides of the file, finds the definition each changed line falls in (a `__global__`, a `__device__` helper and the kernels reaching it, a host launcher and the kernels it launches, a constant and the functions using it) and joins the names back to the map's symbols. Anything it cannot name falls back to the whole file. `CI_SELECTOR_KERNEL_ATTRIBUTION=file` restores the whole-file reading for measurement.
 
-Selecting takes one observation and carries no gate. Dropping carries all of them. For the kernel record the gates are: the row is healthy (every job passed, every shard reported, no dropped records), the file is clearable (compiled into kernels and into no host-only object, so a header `torch_bindings.cpp` includes may select but never drop), the step was selected for nothing but files the record can clear, and the table and map come from the same commit.
+Selecting takes one observation and carries no gate. Dropping carries all of them. For the kernel record the gates are: the row is healthy (every job passed, every shard reported, clean native traces, no dropped records), the file is clearable (compiled into kernels and into no host-only object, so a header `torch_bindings.cpp` includes may select but never drop), the step was selected for nothing but files the record can clear, and the table and map come from the same commit.
 
 ## In CI
 

buildkite/ci_selector/_typos.toml

@@ -4,6 +4,7 @@
 # "literals", not a misspelling of "list". Renaming it collides with the real
 # `literals` already in scope in dispatch.py and keys.py.
 lits = "lits"
+HSA_OVERRIDE_GFX_VERSION = "HSA_OVERRIDE_GFX_VERSION"
 
 [default.extend-words]
 # From "café" in the non-UTF8 fixtures: a latin-1 source file, an undecodable

buildkite/ci_selector/ci_selector/cli.py

@@ -94,12 +94,12 @@ def build_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--kernel-table",
         type=Path,
-        help="kernel table to read (default: coverage-data/kernel_table.json.gz)",
+        help="explicit kernel table override (default: CUDA and ROCm pairs)",
     )
     parser.add_argument(
         "--kernel-symbol-map",
         type=Path,
-        help="kernel symbol map to read (default: coverage-data/kernel_symbol_map.json.gz)",
+        help="explicit kernel symbol map override (default: CUDA and ROCm pairs)",
     )
     return parser
 
@@ -166,13 +166,13 @@ def main(argv: list[str] | None = None) -> int:
     sel = select(state, paths, base=base, head=head)
 
     if args._mode != "codemap":
-        from .coverage.source import fetch_kernel_evidence, fetch_table
+        from .coverage.source import fetch_kernel_records, fetch_table
         from .decide import decide
 
         table = fetch_table(args.table)
         if table.dead_interpreter:
             print(f"coverage: {table.dead_interpreter}", file=sys.stderr)
-        kernels = fetch_kernel_evidence(args.kernel_table, args.kernel_symbol_map)
+        kernels = fetch_kernel_records(args.kernel_table, args.kernel_symbol_map)
         d = decide(state, sel, repo, base, head, table=table, kernels=kernels)
         if d.coverage_note:
             print(f"coverage: {d.coverage_note}", file=sys.stderr)

buildkite/ci_selector/ci_selector/coverage/kernels.py

@@ -12,9 +12,9 @@
 Two files, produced together by a recording build (`recorders/kernrec/README.md`):
 
   kernel_table.json.gz       one row per step: the set of kernel names the
-                             step's processes launched (CUPTI), plus whether
-                             every job passed, every shard reported and no
-                             records were dropped
+                             step's processes launched (CUPTI or ROCProfiler SDK), plus whether
+                             every job passed, every shard reported, traces
+                             completed cleanly and no records were dropped
   kernel_symbol_map.json.gz  per compiled object: its source file, the
                              headers it included, and the kernel entry
                              symbols it defines, read off the objects the
@@ -63,9 +63,10 @@
 from .rules import RowKeys
 
 #: `recorders/kernrec/kernel_table.py::TABLE_VERSION`. A drift test compares the two.
-TABLE_VERSION = 2
+TABLE_VERSION = 3
 #: What `tools/ci/kernel_symbol_map.py` in vLLM writes.
-MAP_VERSION = 1
+MAP_VERSION = 2
+BACKENDS = frozenset({"cuda", "rocm"})
 
 #: Selection rules that carry no changed file: image builds that always run,
 #: preflight force-selects and a pipeline-wide run-all. A step holding one is
@@ -82,11 +83,19 @@ class KernelRow:
     dropped: int
     jobs: int = 0
     processes: int = 0
+    backend: str = "cuda"
+    trace_complete: bool = False
 
     @property
     def usable(self) -> bool:
-        """The same three conditions as `recorders/kernrec/kernel_table.py::usable`."""
-        return self.passed and self.complete and not self.dropped
+        """The same health gates as `recorders/kernrec/kernel_table.py::usable`."""
+        return (
+            self.passed
+            and self.complete
+            and not self.dropped
+            and self.backend in BACKENDS
+            and self.trace_complete
+        )
 
 
 class KernelTable:
@@ -131,6 +140,7 @@ def __init__(
         commit: str = "",
         incomplete: bool = False,
         objects: int = 0,
+        backend: str = "cuda",
         sources: frozenset[str] = frozenset(),
     ):
         self._reach = reach or {}
@@ -144,6 +154,7 @@ def __init__(
         self.commit = commit
         self.incomplete = incomplete
         self.objects = objects
+        self.backend = backend
         # Files an object WITH kernels was compiled from (not its headers):
         # where a host function that launches one can be defined.
         self.sources = sources
@@ -169,6 +180,8 @@ def knows(self, path: str) -> bool:
 
     def why_not_clearable(self, path: str) -> str:
         """Empty when a silence about `path` may drop a step."""
+        if self.incomplete:
+            return "the build's symbol map is incomplete"
         if path in self._unknown:
             return "an object compiled from or including it could not be read"
         if path in self._host:
@@ -223,15 +236,19 @@ def load_table(path: Path) -> KernelTable:
     if payload is None:
         return KernelTable(None, why)
     version = payload.get("version")
-    if version != TABLE_VERSION:
+    if version not in (2, TABLE_VERSION):
         return KernelTable(
             None,
             f"{path} is kernel table version {version!r}, expected {TABLE_VERSION}; "
             "rebuild it from the recordings with recorders/kernrec/kernel_table.py",
         )
     names = payload.get("names")
     blobs = payload.get("rows")
-    if not isinstance(names, list) or not isinstance(blobs, dict):
+    if (
+        not isinstance(names, list)
+        or not all(isinstance(name, str) for name in names)
+        or not isinstance(blobs, dict)
+    ):
         return KernelTable(None, f"{path} has no names list or rows object")
     if not blobs:
         # An empty rows object is a broken table, never "no step launched anything".
@@ -241,6 +258,23 @@ def load_table(path: Path) -> KernelTable:
         for key, b in blobs.items():
             # Indexed, not .get(): a missing health field would default to
             # the healthy value and read a weaker table as stronger.
+            backend = b["backend"] if version == TABLE_VERSION else "cuda"
+            trace_complete = b["trace_complete"] if version == TABLE_VERSION else False
+            if not isinstance(backend, str) or backend not in BACKENDS | {"unknown"}:
+                raise ValueError(f"unknown kernel backend: {backend!r}")
+            if not isinstance(trace_complete, bool):
+                raise ValueError("trace_complete must be a boolean")
+            if any(type(b[field]) is not bool for field in ("passed", "complete")):
+                raise ValueError("passed and complete must be booleans")
+            for field in ("dropped", "jobs", "processes"):
+                value = b.get(field, 0)
+                if type(value) is not int or value < 0:
+                    raise ValueError(f"{field} must be a nonnegative integer")
+            if not isinstance(b["kernels"], list) or any(
+                type(index) is not int or not 0 <= index < len(names)
+                for index in b["kernels"]
+            ):
+                raise ValueError("kernel indices must refer to names")
             rows[key] = KernelRow(
                 key=key,
                 kernels=frozenset(names[i] for i in b["kernels"]),
@@ -249,6 +283,8 @@ def load_table(path: Path) -> KernelTable:
                 dropped=int(b["dropped"]),
                 jobs=int(b.get("jobs", 0)),
                 processes=int(b.get("processes", 0)),
+                backend=backend,
+                trace_complete=trace_complete,
             )
     except (KeyError, TypeError, IndexError, ValueError) as exc:
         return KernelTable(None, f"{path}: unreadable row: {type(exc).__name__}: {exc}")
@@ -267,11 +303,16 @@ def load_symbol_map(path: Path) -> SymbolMap:
     if payload is None:
         return SymbolMap(None, unavailable=why)
     version = payload.get("version")
-    if version != MAP_VERSION:
+    if version not in (1, MAP_VERSION):
         return SymbolMap(
             None,
             unavailable=f"{path} is symbol map version {version!r}, expected {MAP_VERSION}",
         )
+    backend = payload.get("backend") if version == MAP_VERSION else "cuda"
+    if not isinstance(backend, str) or backend not in BACKENDS:
+        return SymbolMap(
+            None, unavailable=f"{path}: unknown kernel backend {backend!r}"
+        )
     # The producer fails closed: any object it could not read after a retry
     # empties the whole map and says why. Nothing to join then.
     if payload.get("reason"):
@@ -310,6 +351,7 @@ def load_symbol_map(path: Path) -> SymbolMap:
         commit=str(payload.get("commit") or ""),
         incomplete=bool(payload.get("incomplete")),
         objects=len(objects),
+        backend=backend,
         sources=frozenset(sources),
     )
 
@@ -321,6 +363,8 @@ class KernelReading:
     added: list[str] = field(default_factory=list)
     dropped: list[str] = field(default_factory=list)
     kept: list[str] = field(default_factory=list)
+    # Positive observations, including steps the map already selected.
+    executes: list[str] = field(default_factory=list)
     reasons: Counter = field(default_factory=Counter)
     #: changed file -> what this record could do with it, for the report
     files: dict[str, str] = field(default_factory=dict)
@@ -433,18 +477,20 @@ def read_pr(
             reading.reasons["file-may-drop"] += 1
 
     def launched(row: KernelRow) -> bool:
-        return any(syms and row.kernels & syms for syms in voting.values())
+        return row.backend == sm.backend and any(
+            syms and row.kernels & syms for syms in voting.values()
+        )
 
     # ADD. Ungated: one launch proves the step runs code from the file, and a
     # failed or incomplete row is still a record of what did run.
     already = set(selection.selected)
-    for step_id in keys.candidates():
-        if step_id in already:
-            continue
+    for step_id in dict.fromkeys([*keys.candidates(), *selection.selected]):
         key = keys.key_for(step_id)
         row = table.row(key) if key else None
         if row is not None and launched(row):
-            reading.added.append(step_id)
+            reading.executes.append(step_id)
+            if step_id not in already:
+                reading.added.append(step_id)
     reading.reasons["row-adds-a-step-the-map-missed"] += len(reading.added)
 
     # DROP. Every gate resolves toward keeping.
@@ -466,6 +512,7 @@ def launched(row: KernelRow) -> bool:
             allow_drops,
             protected,
             held,
+            sm.backend,
             step_id in python_cleared,
         )
         if why_keep:
@@ -495,6 +542,7 @@ def _why_keep(
     allow_drops,
     protected,
     held,
+    backend,
     others_cleared=False,
 ) -> str:
     """The first gate that holds the step, or "" when it may drop."""
@@ -504,6 +552,8 @@ def _why_keep(
     row = table.row(key)
     if row is None:
         return "no-row"
+    if row.backend != backend:
+        return "row-and-map-from-different-backends"
     if launched(row):
         return "row-launched-a-kernel-from-a-changed-file"
     if not row.kernels:

buildkite/ci_selector/ci_selector/coverage/libraries.py

@@ -248,7 +248,7 @@ def read_pr(
             reasons["unmappable-step-id"] += 1
             continue
         if step_id in protected:
-            reasons["held-by-the-python-record"] += 1
+            reasons["held-by-an-execution-record"] += 1
             continue
         unusable = table.unusable(key)
         if unusable is not None:

buildkite/ci_selector/ci_selector/coverage/source.py

@@ -98,3 +98,21 @@ def fetch_kernel_evidence(table_path: Path | None = None, map_path: Path | None
         )
     )
     return KernelEvidence(table, symbol_map)
+
+
+def fetch_kernel_records(table_path: Path | None = None, map_path: Path | None = None):
+    """Load CUDA and ROCm independently, or one explicitly configured pair."""
+    if table_path is not None or map_path is not None:
+        return [fetch_kernel_evidence(table_path, map_path)]
+    rocm_table = Path(
+        os.environ.get(
+            "CI_SELECTOR_ROCM_KERNEL_TABLE", COVERAGE_DIR / "kernel_table.rocm.json.gz"
+        )
+    )
+    rocm_map = Path(
+        os.environ.get(
+            "CI_SELECTOR_ROCM_KERNEL_SYMBOL_MAP",
+            COVERAGE_DIR / "kernel_symbol_map.rocm.json.gz",
+        )
+    )
+    return [fetch_kernel_evidence(), fetch_kernel_evidence(rocm_table, rocm_map)]