diff --git a/metainfer/tasks/opt_GEMM_kernel/README.md b/metainfer/tasks/opt_GEMM_kernel/README.md index 731ee3c2..f634ac66 100644 --- a/metainfer/tasks/opt_GEMM_kernel/README.md +++ b/metainfer/tasks/opt_GEMM_kernel/README.md @@ -1,85 +1,109 @@ # opt_GEMM_kernel -An independent MetaInfer task for arena-style GEMM kernel optimization. It +An independent MetaInfer task for arena-style W8A8 GEMM kernel optimization. It does not import or modify `opt_kernel`, `gen_cpp_infer_framework`, or `gen_infer_framework`. ## Runtime inputs - `initial_submission`: initial HIP challenger and optimization-seed directory. -- `evaluator_bundle`: task-author-provided harness directory containing - `task.yaml` and its correctness and benchmark runners. In the UI this is - called **Harness path**. MetaInfer snapshots it as the system-owned frozen - evaluator before execution. -- `weight_bundle`: task-author-provided `model_weights/` directory containing - `info.json` and one raw `.bin` per tensor. The UI calls it **Weight - directory**. MetaInfer freezes it separately under `system_weights/`, outside - all optimizer-agent workspaces. -- hardware profile selection. The first registered profile is **Hygon K100 / gfx928**. - -The evaluator's `task.yaml::public_contract` owns dtype, layout and ABI, while -its benchmark cases own shapes. These values are parsed once, frozen, supplied -to agents, and displayed read-only; they are not duplicated as manual UI -fields. - -The initial submission includes a constrained `submission.yaml`; it does not -own CMake, compiler or profiler commands. The task-local -`orchestrator/hardware_profiles.yaml` binds the K100 selection to DTK/HIP, -gfx928, CMake + Ninja, `-O3`, and a preferred `hipprof --pmc` route with -rocprofv3/rocprof fallbacks. MetaInfer -resolves the installed executables, -materializes `system_build/{build_profile.json,CMakeLists.txt,build.sh}`, and -freezes the device compiler, host C++ compiler, CMake, Ninja/Make generator, -GPU architecture, fixed flags, and their fingerprint. - -The evaluator bundle is copied into task state before agents run and checked -against a SHA-256 manifest before and after each gate. The optimizer only -receives public notebooks and sanitized feedback. - -The six iteration phases match the C++/Python framework loop exactly: -`A_plan -> B_implement -> C_test -> D_review -> E_perf_test -> F_perf_plan`. -`S_baseline` is a one-time preflight and is not a seventh loop phase. It first -certifies the frozen Triton implementation (correctness, event benchmark, and -PMC) as the iteration-0 Champion, then independently compiles and certifies the -Initial HIP submission with its own correctness, benchmark, PMC, and artifact -directories. Initial HIP replaces Triton only when the existing evaluator and -noise/critical-regression gates accept it. Inside -`C_test`, MetaInfer runs its fixed SystemBuilder and then the harness -correctness command. `E_perf_test` first runs the full frozen event-timed -benchmark and then profiles only three representative public shapes with the -fixed K100 counter groups. The Harness `profile CASE_ID` entrypoint performs -activation generation/quantization, weight loading and copies before its one -candidate GEMM launch, so those preparation costs are not attributed to GEMM. -`D_review` reviews C evidence; `F_perf_plan` analyzes E evidence and prepares -the next optimization. -See `harness/README.md` for the authoring workspace and runtime protocol. - -## Loop +- `evaluator_bundle`: task-author-provided harness containing `task.yaml`, the + correctness runner, and the task-local hipprof suite. The UI calls this + **Harness path**. MetaInfer freezes it before execution. +- `weight_bundle`: task-author-provided `model_weights/` containing `info.json` + and one raw `.bin` per tensor. MetaInfer freezes it under `system_weights/`, + outside every optimizer-agent workspace. +- hardware profile selection. The registered production profile is + **Hygon K100 / gfx928**. + +`task.yaml::public_contract` is the source of truth for dtype, layout, numerics, +and ABI. Its benchmark matrix owns the exact shapes. The UI renders these +values read-only rather than asking the task owner to duplicate them. + +## System-owned execution + +The submission may list source/include paths and allowlisted build options in +`submission.yaml`; it does not own CMake, compiler, architecture, evaluator, or +profiler commands. The K100 hardware profile freezes DTK/HIP, gfx928, CMake + +Ninja, `-O3`, hipprof, all profiler arguments, and a fingerprint of the resolved +tools and protocol. + +The evaluator and weight snapshots are SHA-256 verified at every gate. Agents +receive only the public contract, notebooks, current submission, and sanitized +system evidence. They cannot replace correctness, timing, scoring, or promotion +logic. + +## Performance protocol + +K100 performance latency comes only from the required task-local hipprof trace +suite. For each of the 60 frozen benchmark shapes, candidate and Triton setup, +JIT, allocation, copies, weight preprocessing, packing, workspace initialization, +and synchronization complete before the marked interval. The interval contains +110 steady-state calls: 10 warmup calls followed by 100 measured calls. + +For each logical GEMM call, MetaInfer sums `DurationNs` for every related GPU +dispatch. It then takes the arithmetic mean of the final 100 operator sums. +This is GPU operator time only: host launch API time and synchronization overhead +are excluded. A split-K main kernel plus reduction is therefore one operator +sample containing both GPU dispatch durations. + +Every iteration remeasures the current Champion and candidate in the same +round. Reports retain all raw operator samples and expose mean, median, +standard deviation, CV, and observed range. Results near the noise boundary +trigger a second equal-size hipprof trace for both sides; the decision uses the +arithmetic mean of all raw `DurationNs` operator samples. No shape weighting or +synthetic aggregate latency is used. + +hipprof `--pmc`, `--pmc-read`, and `--pmc-write` run separately. Routine +iterations collect them only for failed diagnostic shapes; a promotable +candidate receives a full-shape PMC archive. They provide +HBM traffic/bandwidth, L2 behavior, VGPR/AGPR/SGPR, LDS, scratch, dispatch and +wave metadata. Occupancy or wave residency is shown only when the profiler +reports a reliable value. PMC replay duration is never latency. Each profiler +pass records its real wall time and has an independent timeout. Missing hipprof, +incomplete cases, unstable dispatch patterns, mismatched protocol fingerprints, +or collection/analyzer failures are infrastructure failures; there is no event +or rocprof timing fallback for K100. + +## Loop and promotion ```text -Certified Triton Champion -> Certified Initial HIP challenger +Certified Triton baseline -> Certified Initial HIP challenger -> A plan -> B implement -> C test -> D review -> E perf test -> F perf plan ``` -Each iteration starts from the persisted HIP Champion source. While Triton is -still Champion, it starts from the independently certified Initial HIP source -because Triton has no editable HIP submission tree. A candidate must pass every -declared correctness and performance case, satisfy the weighted and critical -shape gates, and beat the champion by more than the noise threshold before it -is promoted. - -The task registers its own New Task card and creation form. Its detail page is -kernel-specific: certified hardware/build identity, weighted latency, speedup, -TFLOPS, modelled memory bandwidth, measured memory bandwidth, L2 hit rate, -compute busy, VGPR/LDS pressure, critical-shape regression, per-case profile, -and champion history. Modelled TFLOPS/bandwidth come from frozen evaluator -metadata; hardware counters come only from the frozen system profiler. - -The detail page also provides a live optimization-guidance queue. A task owner can -submit an optimization hypothesis at any time; it is durably delivered to the -next planner or implementer launch and shown as pending/applied in the UI. -Guidance can affect generated candidates but never changes evaluator or -champion gates. - -See `notebooks/02_evaluation_protocol.md` for the evaluator bundle schema and -structured report examples. +`S_baseline` is one-time preflight, not a seventh iteration phase. It certifies +Triton correctness/performance, then independently builds and certifies Initial +HIP. `C_test` runs the system build and frozen correctness command. `E_perf_test` +runs the all-shape hipprof suite and the immutable performance-report gate. + +A candidate must satisfy all of these conditions: + +1. compile and pass every declared correctness case; +2. return one finite positive hipprof operator latency for every benchmark shape; +3. preserve the certified lineage that originally beat Triton on every shape; +4. be below the same-round `champion_ms * (1 - noise_threshold)` on every shape. + +There are no shape weights, critical-shape exceptions, or aggregate score that +can compensate for a losing shape. When Triton remains Champion, the next HIP +iteration still starts from certified Initial HIP because Triton has no editable +HIP submission tree. + +The authoritative performance data is an immutable JSON report referenced by +relative task-state path plus SHA-256. Triton, Initial HIP, every iteration, and +Champion records point to these reports. Cold restart verifies and reloads the +referenced report; iteration scores, timeline fields, and UI summaries are +historical or derived views and never drive promotion. + +The detail page exposes raw per-shape baseline/candidate/Champion latency, +speedup, regression, kernel dispatch breakdown, modeled rates from frozen +metadata, HBM read/write/total bandwidth, L2, registers, LDS/scratch, and +available wave/occupancy evidence. It does not produce a weighted overall score. + +Live task-owner guidance is durable input to the next planner or implementer, +but remains a hypothesis. It cannot alter compilation, correctness, profiler, +all-shape, or Champion gates. + +See `harness/README.md` for harness ownership, +`notebooks/02_evaluation_protocol.md` for report and gate semantics, and +`notebooks/04_profiling.md` for the K100 hipprof route. diff --git a/metainfer/tasks/opt_GEMM_kernel/form.yaml b/metainfer/tasks/opt_GEMM_kernel/form.yaml index 0700e878..256c3654 100644 --- a/metainfer/tasks/opt_GEMM_kernel/form.yaml +++ b/metainfer/tasks/opt_GEMM_kernel/form.yaml @@ -38,7 +38,7 @@ form: select options: - label: "Hygon K100" - description: "DTK/HIP gfx928 with fixed CMake + hipcc and rocprof profiling" + description: "DTK/HIP gfx928 with fixed CMake + hipcc and required hipprof trace/PMC profiling" - key: gpu_arch question: "Compiler target owned by the selected Hygon K100 profile." diff --git a/metainfer/tasks/opt_GEMM_kernel/harness/README.md b/metainfer/tasks/opt_GEMM_kernel/harness/README.md index 3b42cdde..a98baf0e 100644 --- a/metainfer/tasks/opt_GEMM_kernel/harness/README.md +++ b/metainfer/tasks/opt_GEMM_kernel/harness/README.md @@ -1,77 +1,82 @@ # GEMM harness authoring area -This directory is the task-local place for evaluator harnesses. A harness is -provided by the task author; it is not generated or modified by the kernel -optimization agent. +This directory contains task-author-owned evaluator harnesses. A harness is not +generated or modified by an optimization agent. -Select `user_gemm/` in the Web UI's **Harness path** field and the separate -`model_weights/` directory in **Weight directory**. At task start MetaInfer -copies the selected directories to: +Select `user_gemm/` as **Harness path** and the separate `model_weights/` +directory as **Weight directory**. MetaInfer freezes them at task start: ```text /system_evaluator/ /system_weights/ ``` -Both copies are SHA-256 fingerprinted. The evaluator is checked before and -after every command, and the weight directory is outside every agent iteration -workspace. Optimization agents receive only the public contract and sanitized -results, not either private directory. +Both snapshots are SHA-256 fingerprinted. The evaluator is verified around each +system gate, and the weights remain outside agent workspaces. Agents receive the +public contract and sanitized evidence, not private evaluator details. ## Phase ownership ```text -S_baseline MetaInfer build -> harness correctness -> harness benchmark -A_plan agent; no harness execution +S_baseline system build -> correctness -> all-shape hipprof profile +A_plan agent analyzes current source and evidence; no harness edits B_implement agent edits submission/ only -C_test MetaInfer SystemBuilder -> frozen harness correctness command +C_test system build -> frozen correctness command D_review agent reviews compile/correctness evidence -E_perf_test frozen harness benchmark command -> champion decision -F_perf_plan agent analyzes performance and plans the next iteration +E_perf_test frozen all-shape hipprof profile -> Champion decision +F_perf_plan agent analyzes per-shape trace/PMC evidence ``` -`S_baseline` is preflight; the six-phase outer loop is A through F. +`S_baseline` is preflight; the six-phase optimization loop is A through F. +`harness` and `evaluator_bundle` name the same frozen artifact. -Thus `harness` and `evaluator_bundle` refer to the same artifact. The latter is -kept as the requirements/API key for compatibility. +## Required files and ownership -## Required files +Every selectable harness contains `task.yaml`. It defines the public contract, +correctness cases, benchmark shapes, frozen hipprof protocol, correctness +command, and profile entry point. Correctness writes a JSON object to +`METAINFER_REPORT_PATH` and returns zero only after the reference checks pass. +Performance reports are generated by the system-owned profiler runner from the +frozen task-local hipprof suite; agents do not supply a benchmark command. -Every selectable harness directory must contain `task.yaml`. Its commands must -write a JSON object to `METAINFER_REPORT_PATH` and return zero only when the -phase completed normally and its report is valid. - -MetaInfer supplies these environment variables: +MetaInfer supplies the relevant environment variables: - `METAINFER_EVALUATOR_BUNDLE`: frozen harness directory. - `METAINFER_SUBMISSION_DIR`: source submission being evaluated. -- `METAINFER_BUILD_ARTIFACT_DIR`: system-built candidate artifact directory. -- `METAINFER_REPORT_PATH`: required JSON output path. -- `METAINFER_EVALUATION_PHASE`: `correctness` or `benchmark`. +- `METAINFER_BUILD_ARTIFACT_DIR`: system-built candidate artifacts. +- `METAINFER_REPORT_PATH`: required system report path. +- `METAINFER_EVALUATION_PHASE`: current system gate. - `METAINFER_EVALUATION_ROLE`: `baseline` or `candidate`. - `METAINFER_BUILD_FINGERPRINT`: frozen compiler/build identity. -- `METAINFER_BENCHMARK_PROTOCOL`: frozen JSON timing protocol. -- `METAINFER_WEIGHT_BUNDLE`: frozen directory containing `info.json` and the - separate tensor `.bin` files. -- `METAINFER_WEIGHT_SHA256`: fingerprint of that frozen weight directory. - -The harness should locate and load the candidate shared library from -`METAINFER_BUILD_ARTIFACT_DIR`. Do not compile the candidate itself: CMake, -hipcc/nvcc, target architecture and candidate flags are the first internal -gate of `C_test` and remain owned by MetaInfer. - -## Trust rules - -- Put CPU/PyTorch references, input generation, tolerances and case definitions - in the harness. -- Include all correctness cases in the JSON report, including private cases. - MetaInfer removes private details before feedback reaches an agent. -- Benchmark only the operation covered by the public ABI. Exclude allocation, - host/device copies and process startup from `latency_ms`. -- Use deterministic inputs, GPU-event timing, warmup and repeated samples. -- Never report success before the reference comparison actually passes. -- Keep harness build products outside this source directory so the frozen - bundle digest remains stable. - -`user_gemm/evaluate_native.cpp` is the concrete W8A8 runner for the supplied tensor -metadata. Its README documents the TP4/TP8 slicing and concatenation rules. +- `METAINFER_BENCHMARK_PROTOCOL`: frozen hipprof timing protocol JSON. +- `METAINFER_WEIGHT_BUNDLE`: frozen tensor directory. +- `METAINFER_WEIGHT_SHA256`: frozen weight fingerprint. + +The harness loads the candidate library from +`METAINFER_BUILD_ARTIFACT_DIR`. It must not compile the candidate or choose a +compiler, GPU architecture, profiler command, counter group, or timing fallback. +Those are owned and fingerprinted by MetaInfer. + +## Correctness and performance trust rules + +- Put independent references, deterministic input generation, tolerances, and + case definitions in the harness. +- Include every correctness case in the report. MetaInfer sanitizes private + details before agent feedback. +- Complete activation preparation, JIT, allocation, copies, packing, workspace + initialization, and synchronization before the marked profiling interval. +- Put only repeated steady-state ABI calls in the marked interval. +- Use hipprof trace `DurationNs` as the sole K100 latency source. Sum every GPU + dispatch belonging to one logical call, then average the frozen final samples. +- Validate exact call count and a stable final dispatch pattern for every shape. +- Collect PMC/read/write in separate replay passes. Use counters only for + traffic, cache, resource, and reliably reported occupancy/wave diagnostics; + never use replay duration as latency. +- Do not attach weights or criticality to benchmark shapes. Every shape is an + independent hard gate. +- Keep harness build/profile products outside this source directory so the + frozen digest remains stable. + +`user_gemm/evaluate_native.cpp` is the concrete correctness runner for the +supplied W8A8 tensors. `user_gemm/README.md` documents TP4/TP8 derivation and the +exact task-local profile protocol. diff --git a/metainfer/tasks/opt_GEMM_kernel/harness/user_gemm/README.md b/metainfer/tasks/opt_GEMM_kernel/harness/user_gemm/README.md index 79759769..2ab5267b 100644 --- a/metainfer/tasks/opt_GEMM_kernel/harness/user_gemm/README.md +++ b/metainfer/tasks/opt_GEMM_kernel/harness/user_gemm/README.md @@ -1,19 +1,18 @@ # Scaled W8A8 GEMM harness -This directory is a complete, task-author-owned evaluator for the supplied -DeepSeek-style W8A8 weights. Select this directory as **Harness path**, and -select the separate `model_weights/` directory as **Weight directory**. -MetaInfer snapshots both under task state before the baseline runs. Optimization -agents receive the public ABI and shapes, but cannot edit either snapshot. +This directory is the task-author-owned evaluator for the supplied +DeepSeek-style W8A8 weights. Select it as **Harness path** and select the +separate `model_weights/` directory as **Weight directory**. MetaInfer freezes +both before baseline certification. Agents receive the public ABI, shapes, and +sanitized evidence but cannot edit either snapshot. -The matrix-multiply baseline is kept in the separate initial submission at -`../../initial_submissions/myGEMM_kernel/`. Select that directory as **Kernel -path**. Combined demo code containing `main()`, allocation, testing, and timing -does not live in the frozen Harness or candidate submission. +The editable HIP seed is under `../../initial_submissions/myGEMM_kernel/` and is +selected as **Kernel path**. Allocation, references, testing, and profiling +control do not belong in a candidate submission. -## Required weight directory +## Required weights and TP rank 0 derivation -`model_weights/` must contain `info.json` plus these separate raw files: +`model_weights/` contains `info.json` and one raw file for each tensor/scale: ```text q_proj_a.bin q_proj_a_scale.bin @@ -25,45 +24,35 @@ moe_w2.bin moe_w2_scale.bin moe_w3.bin moe_w3_scale.bin ``` -`evaluate_native.cpp` checks every shape and exact file length against the -metadata supplied for this task. It does not assume a concatenated binary or -byte offsets. - -## Weight derivation for TP rank 0 +`evaluate_native.cpp` validates every filename, dtype, shape, and exact byte +length. It does not assume concatenated binaries or hidden offsets. - `wqkv_a`: concatenate `q_proj_a` and `kv_proj` on N; unchanged for TP4/TP8. -- `wq_b`: take the first `32768 / TP` columns of `q_proj_b` and its scale. -- `wo_b`: take the first `8192 / TP` rows of `o_proj`; output scale is unchanged. -- `shared_gate_up_proj`: take the first `2048 / TP` columns from each of - `moe_w1` and `moe_w3`, then concatenate them and their scales on N. -- `shared_down_proj`: take the first `2048 / TP` rows of `moe_w2`; output scale - is unchanged. +- `wq_b`: first `32768 / TP` columns of `q_proj_b` and its scale. +- `wo_b`: first `8192 / TP` rows of `o_proj`; output scale is unchanged. +- `shared_gate_up_proj`: first `2048 / TP` columns from each of `moe_w1` and + `moe_w3`, then concatenate on N. +- `shared_down_proj`: first `2048 / TP` rows of `moe_w2`; output scale is + unchanged. -All loading, slicing, concatenation and host-to-device copies occur outside the -timed interval. `indexer.wq_b` is intentionally excluded until its independent -weight tensor and scale are supplied. +Loading, slicing, concatenation, packing, and host-to-device copies occur before +the marked interval. `indexer.wq_b` remains excluded until its independent +weight and scale are supplied. -## Activation and timed scope +## Activation, correctness, and ABI -For each case the harness deterministically generates BF16 `A[M,K]`, then does -per-row symmetric quantization: +The harness deterministically generates BF16 `A[M,K]`, then performs per-row +symmetric quantization: ```text A_scale[m] = max(abs(A[m,:])) / 127 A_int8 = clamp(round(A / A_scale), -127, 127) ``` -The candidate receives `A_int8`, `W_int8`, `A_scale`, and `W_scale`. GPU events -measure only `launch_w8a8_gemm(...)`; activation quantization, allocation, -weight preprocessing and copies are excluded. - -Correctness checks the complete result against a frozen, independent GPU INT32 -reference kernel and also recomputes deterministic sentinel points with CPU -INT64 accumulation. -Benchmarking covers TP4/TP8 and `M = 1,2,4,8,16,4096` with 10 warmups and 100 -GPU-event samples per case. - -## Candidate ABI +The candidate receives prepared `A_int8`, `W_int8`, `A_scale`, and `W_scale` and +produces only the BF16 result. Correctness compares the complete output with a +frozen independent GPU INT32 reference and recomputes deterministic sentinel +points with CPU INT64 accumulation. ```cpp extern "C" int launch_w8a8_gemm( @@ -78,7 +67,37 @@ extern "C" int launch_w8a8_gemm( void* stream); ``` -Return zero after enqueueing work on the supplied stream. The shared library -and frozen native harness executable are built together by MetaInfer's fixed -CMake/hipcc or CMake/nvcc route. The harness then loads the candidate library -from `METAINFER_BUILD_ARTIFACT_DIR`. +Return zero after enqueueing work on the supplied stream. MetaInfer owns the +fixed CMake/hipcc build and loads the resulting library from +`METAINFER_BUILD_ARTIFACT_DIR`. + +## Task-local hipprof performance protocol + +K100 latency is collected only by the frozen scripts in this directory: + +```bash +python3 run_hipprof_suite.py --output-dir "$METAINFER_REPORT_DIR/hipprof-suite" +python3 analyze_hipprof_suite.py "$METAINFER_REPORT_DIR/hipprof-suite" +``` + +The system runner launches the suite with its actual Python interpreter, the +frozen candidate artifact, frozen weights, and frozen benchmark protocol. The +matrix contains TP4/TP8 workloads at `M = 1,2,4,8,16,4096`, for 60 shapes total. +Candidate and Triton setup, JIT, allocation, copies, packing, workspace setup, +and synchronization finish before each marked interval. + +Each trace interval contains 110 steady-state logical calls. The first 10 are +warmup and the final 100 are measured. For every call, the analyzer sums +`DurationNs` for all GEMM GPU dispatches, including split-K and reduction, then +takes the arithmetic mean of the final 100 sums. It verifies the exact call +count and stable final dispatch pattern. Host launch API and synchronization +time are outside this GPU operator latency. + +Separate `--pmc`, `--pmc-read`, and `--pmc-write` passes provide HBM read/write +bytes and bandwidth, L2 hit behavior, VGPR/AGPR/SGPR, LDS, scratch, dispatch, +workgroup, and wave metadata. Occupancy remains unavailable unless hipprof +reports a reliable value. PMC replay duration is never used as latency. + +Every one of the 60 shapes is a hard performance gate. A candidate must be +strictly faster than frozen Triton on each shape and cross the current Champion +noise threshold on each shape; no weight or aggregate average can hide a loss. diff --git a/metainfer/tasks/opt_GEMM_kernel/harness/user_gemm/analyze_hipprof_suite.py b/metainfer/tasks/opt_GEMM_kernel/harness/user_gemm/analyze_hipprof_suite.py new file mode 100644 index 00000000..0b92ee17 --- /dev/null +++ b/metainfer/tasks/opt_GEMM_kernel/harness/user_gemm/analyze_hipprof_suite.py @@ -0,0 +1,354 @@ +#!/usr/bin/env python3 +"""Analyze task-local hipprof traces and PMC into per-shape core metrics.""" + +from __future__ import annotations + +import argparse +import csv +import json +import math +import re +import sqlite3 +import statistics +import time +from pathlib import Path +from typing import Any + + +HARDWARE_READ_GBS = 608.357 +INT8_PEAK_TOPS = 123.310 + + +def _manifest(root: Path, label: str) -> dict[str, Any]: + return json.loads( + (root / f"{label}-harness.json").read_text(encoding="utf-8")) + + +def _trace_rows(path: Path) -> list[dict[str, Any]]: + with sqlite3.connect(path) as db: + names = dict(db.execute( + "SELECT CAST(STR_ID AS TEXT), STR_NAME FROM STR_TABLE WHERE TYPE=6")) + tables = [row[0] for row in db.execute( + "SELECT name FROM sqlite_master WHERE type='table'") + if row[0].startswith("HIPOPS_")] + rows = [] + for table in tables: + for begin, duration, name in db.execute( + f'SELECT BeginNs, DurationNs, CAST(Name AS TEXT) FROM "{table}"'): + rows.append({ + "begin_ns": int(begin), "duration_ns": int(duration), + "kernel_name": names.get(name, name), + }) + return sorted(rows, key=lambda row: row["begin_ns"]) + + +def _number(value: Any) -> float: + try: + return float(value or 0) + except (TypeError, ValueError): + return 0.0 + + +def _sum_indexed(row: dict[str, str], base: str) -> int: + pattern = re.compile(rf"^{re.escape(base)}\[(\d+)\]$") + return sum( + int(_number(value)) for key, value in row.items() + if pattern.match(key)) + + +def _read_bytes(row: dict[str, str]) -> int: + total = 0 + for prefix in ("TCC_EA", "TCC_EA1"): + req = _sum_indexed(row, f"{prefix}_RDREQ") + req32 = _sum_indexed(row, f"{prefix}_RDREQ_32B") + total += req32 * 32 + (req - req32) * 64 + return total + + +def _write_bytes(row: dict[str, str]) -> int: + total = 0 + for prefix in ("TCC_EA", "TCC_EA1"): + req = _sum_indexed(row, f"{prefix}_WRREQ") + req64 = _sum_indexed(row, f"{prefix}_WRREQ_64B") + total += (req - req64) * 32 + req64 * 64 + return total + + +def _counter_rows(path: Path) -> list[dict[str, str]]: + with path.open(newline="", encoding="utf-8") as stream: + return list(csv.DictReader(stream)) + + +def _begin(row: dict[str, str]) -> int: + for key in ("BeginNs", "Begin_Ns", "StartNs", "Start_Timestamp"): + if row.get(key): + return int(_number(row[key])) + raise RuntimeError("hipprof CSV has no begin timestamp column") + + +def _in_case(rows: list[Any], case: dict[str, Any], begin_fn) -> list[Any]: + lo = int(case["host_monotonic_begin_ns"]) + hi = int(case["host_monotonic_end_ns"]) + return [row for row in rows if lo <= begin_fn(row) <= hi] + + +def _trace_case_times( + rows: list[dict[str, Any]], + case: dict[str, Any], + calls: int, + samples: int, + epoch_offset_ns: int = 0, +) -> tuple[float, dict[str, float], int, list[float]]: + if "host_epoch_begin_ns" in case: + lo = int(case["host_epoch_begin_ns"]) + hi = int(case["host_epoch_end_ns"]) + else: + # Compatibility for collections produced before the evaluator wrote + # realtime boundaries. hipprof trace DB timestamps are CLOCK_REALTIME, + # while PMC CSV and the old manifest use CLOCK_MONOTONIC. + lo = int(case["host_monotonic_begin_ns"]) + epoch_offset_ns + hi = int(case["host_monotonic_end_ns"]) + epoch_offset_ns + selected = [row for row in rows if lo <= row["begin_ns"] <= hi] + if not selected or len(selected) % calls: + raise RuntimeError( + f"{case['id']}: {len(selected)} trace dispatches not divisible by {calls}") + if samples < 1 or samples > calls: + raise RuntimeError(f"{case['id']}: invalid trace sample count {samples}") + dispatches_per_call = len(selected) // calls + patterns = [ + tuple(row["kernel_name"] for row in selected[ + index * dispatches_per_call:(index + 1) * dispatches_per_call]) + for index in range(calls) + ] + if any(pattern != patterns[-1] for pattern in patterns[-samples:]): + raise RuntimeError(f"{case['id']}: unstable measured dispatch pattern") + + operator_values = [] + kernel_contributions: dict[str, list[float]] = {} + for index in range(calls - samples, calls): + group = selected[index * dispatches_per_call:(index + 1) * dispatches_per_call] + operator_values.append(sum(row["duration_ns"] for row in group) / 1000.0) + per_call: dict[str, float] = {} + for row in group: + name = row["kernel_name"] + per_call[name] = per_call.get(name, 0.0) + row["duration_ns"] / 1000.0 + for name, contribution in per_call.items(): + kernel_contributions.setdefault(name, []).append(contribution) + return statistics.fmean(operator_values), { + name: statistics.fmean(items) + for name, items in kernel_contributions.items() + }, dispatches_per_call, operator_values + + +def _aggregate_counters( + rows: list[dict[str, str]], case: dict[str, Any], calls: int, +) -> dict[str, Any]: + selected = _in_case(rows, case, _begin) + if not selected: + raise RuntimeError(f"{case['id']}: no PMC dispatches in host interval") + if calls < 1 or len(selected) % calls: + raise RuntimeError( + f"{case['id']}: {len(selected)} PMC dispatches not divisible by {calls}" + ) + # DTK hipprof --pmc-type 3 performs six hardware replay passes internally, + # then exports their counter groups on one merged CSV row per original + # dispatch. Do not apply the older values[index * 6 + 5] rule to this + # merged format: it would discard five real operator dispatches. + # Sum traffic/caches across every original operator dispatch, and use the + # longest dispatch for per-kernel resource metadata. + main = max(selected, key=lambda row: _number( + row.get("DurationNs") or row.get("DispatchNs") or 0)) + hits = sum(_sum_indexed(row, "TCC_HIT") for row in selected) + misses = sum(_sum_indexed(row, "TCC_MISS") for row in selected) + return { + "dispatch_count": len(selected) // calls, + "hbm_read_bytes": sum(_read_bytes(row) for row in selected) / calls, + "hbm_write_bytes": sum(_write_bytes(row) for row in selected) / calls, + "l2_hit_pct": 100.0 * hits / (hits + misses) if hits + misses else math.nan, + "vgpr": int(_number(main.get("arch_vgpr") or main.get("VGPR_Count"))), + "agpr": int(_number(main.get("accum_vgpr") or main.get("AGPR_Count"))), + "sgpr": int(_number(main.get("sgpr") or main.get("SGPR_Count"))), + "lds_bytes": int(_number(main.get("lds") or main.get("LDS_Block_Size"))), + "scratch_bytes": int(_number(main.get("scr") or main.get("Scratch_Size"))), + "grid_size": int(_number(main.get("grd") or main.get("GridSize"))), + "workgroup_size": int(_number(main.get("wgr") or main.get("WorkgroupSize"))), + "wave_size": int(_number(main.get("wave_size") or main.get("WaveSize"))), + "waves_per_workgroup": ( + math.ceil( + _number(main.get("wgr") or main.get("WorkgroupSize")) + / _number(main.get("wave_size") or main.get("WaveSize")) + ) + if _number(main.get("wgr") or main.get("WorkgroupSize")) > 0 + and _number(main.get("wave_size") or main.get("WaveSize")) > 0 + else None + ), + "occupancy_pct": None, + "main_kernel": main.get("KernelName", ""), + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("profile_dir", type=Path) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + root = args.profile_dir.resolve() + collection = json.loads( + (root / "collection.json").read_text(encoding="utf-8")) + collection_mode = str(collection.get("passes") or "full") + if collection_mode == "diagnostic": + raise RuntimeError("diagnostic collection must be analyzed with a trace collection") + rows = [] + for impl in collection["implementations"]: + trace_manifest = _manifest(root, f"{impl}-trace") + trace = _trace_rows(root / f"{impl}-trace.db") + epoch_offset_ns = 0 + if trace and trace_manifest["cases"] and "host_epoch_begin_ns" not in trace_manifest["cases"][0]: + # CLOCK_REALTIME - CLOCK_MONOTONIC is stable for the lifetime of + # a boot. This accurately translates legacy manifests collected + # on this worker; using the first trace row is wrong because it + # includes each case's unmarked Triton warm-up launch. + epoch_offset_ns = time.time_ns() - time.monotonic_ns() + has_pmc = collection_mode == "full" + manifests = {} + counter_data = {} + if has_pmc: + manifests = { + label: _manifest(root, f"{impl}-{label}") + for label in ("pmc", "read", "write") + } + counter_data = { + label: _counter_rows(root / f"{impl}-{label}.csv") + for label in ("pmc", "read", "write") + } + expected_trace_calls = int(collection["trace_calls"]) + if int(trace_manifest.get("calls_per_case") or 0) != expected_trace_calls: + raise RuntimeError( + f"{impl}: trace call count differs from frozen collection protocol" + ) + expected_pmc_calls = int(collection["pmc_calls"]) + for label, manifest in manifests.items(): + if int(manifest.get("calls_per_case") or 0) != expected_pmc_calls: + raise RuntimeError( + f"{impl}: {label} call count differs from frozen collection protocol" + ) + case_maps = { + label: {str(item["id"]): item for item in manifest["cases"]} + for label, manifest in manifests.items() + } + trace_ids = [str(item["id"]) for item in trace_manifest["cases"]] + for label, case_map in case_maps.items(): + if set(case_map) != set(trace_ids): + raise RuntimeError( + f"{impl}: {label} manifest cases differ from trace manifest" + ) + + samples = int(collection["samples"]) + for case in trace_manifest["cases"]: + case_id = str(case["id"]) + operator_us, kernels, trace_dispatches, operator_samples = _trace_case_times( + trace, + case, + int(trace_manifest["calls_per_case"]), + samples, + epoch_offset_ns, + ) + if not kernels: + raise RuntimeError(f"{case_id}: no GPU kernel dispatch") + timed_kernel = max(kernels.items(), key=lambda item: item[1])[0] + aligned = {label: case_maps[label][case_id] for label in case_maps} + for label, manifest_case in aligned.items(): + if any( + int(manifest_case[key]) != int(case[key]) + for key in ("m", "n", "k") + ): + raise RuntimeError( + f"{case_id}: {label} manifest shape differs from trace" + ) + empty_meta = { + "dispatch_count": 0, "hbm_read_bytes": 0, "hbm_write_bytes": 0, + "l2_hit_pct": math.nan, "vgpr": 0, "agpr": 0, "sgpr": 0, + "lds_bytes": 0, "scratch_bytes": 0, "grid_size": 0, + "workgroup_size": 0, "wave_size": 0, + "waves_per_workgroup": None, "occupancy_pct": None, + "main_kernel": "", + } + meta = dict(empty_meta) + read_meta = dict(empty_meta) + write_meta = dict(empty_meta) + if has_pmc: + meta = _aggregate_counters( + counter_data["pmc"], aligned["pmc"], + int(manifests["pmc"]["calls_per_case"]) + ) + read_meta = _aggregate_counters( + counter_data["read"], aligned["read"], + int(manifests["read"]["calls_per_case"]) + ) + write_meta = _aggregate_counters( + counter_data["write"], aligned["write"], + int(manifests["write"]["calls_per_case"]) + ) + m, n, k = (int(case[key]) for key in ("m", "n", "k")) + seconds = operator_us * 1e-6 + rd = read_meta["hbm_read_bytes"] + wr = write_meta["hbm_write_bytes"] + rows.append({ + "case_id": case_id, "implementation": impl, + "M": m, "N": n, "K": k, + "operator_mean_us": operator_us, + "operator_median_us": statistics.median(operator_samples), + "operator_stddev_us": ( + statistics.stdev(operator_samples) + if len(operator_samples) > 1 else 0.0 + ), + "operator_cv": ( + statistics.stdev(operator_samples) / operator_us + if len(operator_samples) > 1 and operator_us > 0 else 0.0 + ), + "operator_min_us": min(operator_samples), + "operator_max_us": max(operator_samples), + "operator_samples_us": operator_samples, + "trace_dispatches_per_call": trace_dispatches, + "timed_kernel": timed_kernel, + "effective_int8_tops": 2 * m * n * k / seconds / 1e12, + "hbm_read_bytes": int(rd), "hbm_write_bytes": int(wr), + "hbm_read_gbs": rd / seconds / 1e9, + "hbm_write_gbs": wr / seconds / 1e9, + "hbm_total_gbs": (rd + wr) / seconds / 1e9, + "hbm_read_attainment_pct": + 100.0 * rd / seconds / 1e9 / HARDWARE_READ_GBS, + "l2_hit_pct": meta["l2_hit_pct"], + "vgpr": meta["vgpr"], "agpr": meta["agpr"], + "sgpr": meta["sgpr"], "lds_bytes": meta["lds_bytes"], + "scratch_bytes": meta["scratch_bytes"], + "grid_size": meta["grid_size"], + "workgroup_size": meta["workgroup_size"], + "wave_size": meta["wave_size"], + "waves_per_workgroup": meta["waves_per_workgroup"], + "occupancy_pct": meta["occupancy_pct"], + "pmc_dispatch_count": meta["dispatch_count"], + "main_kernel": meta["main_kernel"] or timed_kernel, + "trace_kernel_means_json": json.dumps(kernels, sort_keys=True), + "hardware_read_peak_gbs": HARDWARE_READ_GBS, + "int8_compute_peak_tops": INT8_PEAK_TOPS, + }) + output = (args.output or root / "metrics.csv").resolve() + with output.open("w", newline="", encoding="utf-8") as stream: + writer = csv.DictWriter(stream, fieldnames=list(rows[0])) + writer.writeheader(); writer.writerows(rows) + (root / "metrics.json").write_text(json.dumps({ + "timing": ( + "hipprof arithmetic mean after frozen warmup; " + "each operator sample sums all GPU kernel dispatch durations" + ), + "hbm_peak_read_gbs": HARDWARE_READ_GBS, + "collection_mode": collection_mode, + "pass_records": collection.get("pass_records") or [], + "rows": rows, + }, indent=2) + "\n", encoding="utf-8") + print(f"wrote {output}") + + +if __name__ == "__main__": + main() diff --git a/metainfer/tasks/opt_GEMM_kernel/harness/user_gemm/evaluate.py b/metainfer/tasks/opt_GEMM_kernel/harness/user_gemm/evaluate.py index dd302d9d..f7441302 100644 --- a/metainfer/tasks/opt_GEMM_kernel/harness/user_gemm/evaluate.py +++ b/metainfer/tasks/opt_GEMM_kernel/harness/user_gemm/evaluate.py @@ -1,15 +1,13 @@ #!/usr/bin/env python3 -"""Python evaluator for opt_GEMM_kernel — Triton as reference and baseline. +"""Frozen correctness harness and hipprof workload driver. -Replaces evaluate_native.cpp. Uses Triton matmul_int8 as the correctness -reference AND as the performance baseline, so the MetaInfer optimization -loop chases Triton-level (MFMA) throughput. +Triton is the independent correctness reference and the frozen performance +baseline. Performance measurements are produced only by task-local hipprof +trace collection around ``profile-batch`` steady-state GPU dispatches. Phases: - correctness – candidate vs Triton, per-element comparison - benchmark – GPU-event timed measurement (Triton for baseline role, - candidate .so for candidate role) - profile – single candidate launch (wrapped by rocprof) + correctness – candidate vs Triton, per-element comparison + profile-batch – all public cases, repeated steady-state calls for hipprof """ from __future__ import annotations @@ -365,158 +363,72 @@ def _run_triton_correctness_case( # ═══════════════════════════════════════════════════════════════════════════════ -# benchmark +# hipprof workload # ═══════════════════════════════════════════════════════════════════════════════ -def _benchmark_case_triton( - weights: WeightStore, - case: Case, - device: torch.device, - warmup: int, - samples: int, -) -> Dict[str, Any]: - """Benchmark Triton matmul_int8 with GPU events.""" - a_bf16, a_int8, a_scale = _generate_activation(case) - w_int8_np, w_scale_np = weights.derive(case) - - a_int8_dev = a_int8.to(device) - a_scale_dev = a_scale.to(device) - w_int8_dev = torch.from_numpy(w_int8_np).to(device) - w_scale_dev = torch.from_numpy(w_scale_np).to(device) - - # Warmup - for _ in range(warmup): - matmul_int8(a_int8_dev, a_scale_dev, w_int8_dev, w_scale_dev, torch.bfloat16, None) - torch.cuda.synchronize() - - values = [] - for _ in range(samples): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - matmul_int8(a_int8_dev, a_scale_dev, w_int8_dev, w_scale_dev, torch.bfloat16, None) - end.record() - torch.cuda.synchronize() - values.append(start.elapsed_time(end)) - - values.sort() - latency = values[len(values) // 2] - flops = 2.0 * case.m * case.n * case.k - # rough byte count: A(bf16)+W(int8)+scales+output(bf16) - nbytes = (case.m * case.k * 2 + case.k * case.n * 1 - + case.m * 4 + case.n * 4 + case.m * case.n * 2) - - return { - "id": case.id, - "latency_ms": latency, - "min_ms": values[0], - "max_ms": values[-1], - "tops": flops / (latency * 1e9), - "bandwidth_gbps": nbytes / (latency * 1e6), - } - - -def _benchmark_case_candidate( +def _profile_batch_case_candidate( candidate: Candidate, weights: WeightStore, case: Case, device: torch.device, - warmup: int, - samples: int, -) -> Dict[str, Any]: - """Benchmark candidate .so with GPU events.""" - a_bf16, a_int8, a_scale = _generate_activation(case) + calls: int, +) -> Tuple[int, int, int, int]: + """Prepare once, then enqueue exactly ``calls`` candidate invocations.""" + _, a_int8, a_scale = _generate_activation(case) w_int8_np, w_scale_np = weights.derive(case) - - a_int8_dev = a_int8.to(device) - a_scale_dev = a_scale.to(device) - w_int8_dev = torch.from_numpy(w_int8_np).to(device) - w_scale_dev = torch.from_numpy(w_scale_np).to(device) + a_dev = a_int8.to(device) + as_dev = a_scale.to(device) + w_dev = torch.from_numpy(w_int8_np).to(device) + ws_dev = torch.from_numpy(w_scale_np).to(device) y = torch.empty((case.m, case.n), dtype=torch.bfloat16, device=device) - - # Warmup - for _ in range(warmup): - ret = candidate.launch(a_int8_dev, w_int8_dev, a_scale_dev, w_scale_dev, y) - if ret != 0: - raise RuntimeError(f"candidate returned non-zero for {case.id}") + # Complete one-time weight packing, workspace allocation, and lazy runtime + # setup before the profiler's marked steady-state interval. + ret = candidate.launch(a_dev, w_dev, as_dev, ws_dev, y) + if ret != 0: + raise RuntimeError(f"candidate returned non-zero for {case.id}") torch.cuda.synchronize() - - values = [] - for _ in range(samples): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - ret = candidate.launch(a_int8_dev, w_int8_dev, a_scale_dev, w_scale_dev, y) + begin_ns = time.monotonic_ns() + begin_epoch_ns = time.time_ns() + print( + f"PROFILE_GROUP,candidate,{case.id},{case.m},{case.n},{case.k}," + f"calls={calls}", flush=True, + ) + for _ in range(calls): + ret = candidate.launch(a_dev, w_dev, as_dev, ws_dev, y) if ret != 0: raise RuntimeError(f"candidate returned non-zero for {case.id}") - end.record() - torch.cuda.synchronize() - values.append(start.elapsed_time(end)) - - values.sort() - latency = values[len(values) // 2] - flops = 2.0 * case.m * case.n * case.k - nbytes = (case.m * case.k * 2 + case.k * case.n * 1 - + case.m * 4 + case.n * 4 + case.m * case.n * 2) - - return { - "id": case.id, - "latency_ms": latency, - "min_ms": values[0], - "max_ms": values[-1], - "tops": flops / (latency * 1e9), - "bandwidth_gbps": nbytes / (latency * 1e6), - } - - -# ═══════════════════════════════════════════════════════════════════════════════ -# profile -# ═══════════════════════════════════════════════════════════════════════════════ - - -def _profile_case( - candidate: Candidate, - weights: WeightStore, - case: Case, - device: torch.device, -) -> None: - """Single candidate launch for rocprof capture.""" - a_bf16, a_int8, a_scale = _generate_activation(case) - w_int8_np, w_scale_np = weights.derive(case) - - a_int8_dev = a_int8.to(device) - a_scale_dev = a_scale.to(device) - w_int8_dev = torch.from_numpy(w_int8_np).to(device) - w_scale_dev = torch.from_numpy(w_scale_np).to(device) - y = torch.empty((case.m, case.n), dtype=torch.bfloat16, device=device) - - torch.cuda.synchronize() - ret = candidate.launch(a_int8_dev, w_int8_dev, a_scale_dev, w_scale_dev, y) - if ret != 0: - raise RuntimeError(f"candidate returned non-zero for {case.id}") torch.cuda.synchronize() + return begin_ns, time.monotonic_ns(), begin_epoch_ns, time.time_ns() -def _profile_case_triton( +def _profile_batch_case_triton( weights: WeightStore, case: Case, device: torch.device, -) -> None: - """Warm up Triton JIT, then launch exactly one profiled invocation.""" + calls: int, +) -> Tuple[int, int, int, int]: + """JIT before the marked group, then enqueue fixed Triton invocations.""" _, a_int8, a_scale = _generate_activation(case) w_int8_np, w_scale_np = weights.derive(case) - a_int8_dev = a_int8.to(device) - a_scale_dev = a_scale.to(device) - w_int8_dev = torch.from_numpy(w_int8_np).to(device) - w_scale_dev = torch.from_numpy(w_scale_np).to(device) - # Triton's disk cache is populated by certification benchmark. This is - # the single matmul invocation observed by hipprof in this process. - matmul_int8( - a_int8_dev, a_scale_dev, w_int8_dev, w_scale_dev, - torch.bfloat16, None, + a_dev = a_int8.to(device) + as_dev = a_scale.to(device) + w_dev = torch.from_numpy(w_int8_np).to(device) + ws_dev = torch.from_numpy(w_scale_np).to(device) + # Force JIT/allocation before the group marker. The analyzer uses the + # manifest and final repeated core launches, never this preparation call. + matmul_int8(a_dev, as_dev, w_dev, ws_dev, torch.bfloat16, None) + torch.cuda.synchronize() + begin_ns = time.monotonic_ns() + begin_epoch_ns = time.time_ns() + print( + f"PROFILE_GROUP,triton,{case.id},{case.m},{case.n},{case.k}," + f"calls={calls}", flush=True, ) + for _ in range(calls): + matmul_int8(a_dev, as_dev, w_dev, ws_dev, torch.bfloat16, None) torch.cuda.synchronize() + return begin_ns, time.monotonic_ns(), begin_epoch_ns, time.time_ns() # ═══════════════════════════════════════════════════════════════════════════════ @@ -526,11 +438,11 @@ def _profile_case_triton( def main() -> None: phase = sys.argv[1] - is_profile = phase == "profile" - is_eval = phase in ("correctness", "benchmark") - if not is_profile and not is_eval: + is_profile_batch = phase == "profile-batch" + is_correctness = phase == "correctness" + if not is_profile_batch and not is_correctness: raise RuntimeError( - "usage: evaluate.py correctness|benchmark|profile CASE_ID" + "usage: evaluate.py correctness|profile-batch candidate|triton CALLS" ) report_path = Path(_env("METAINFER_REPORT_PATH")) @@ -544,24 +456,54 @@ def main() -> None: device = torch.device("cuda:0") weights = WeightStore(weight_root) - candidate = None if role == "baseline" else Candidate(artifact_dir) - - if is_profile: - case_id = sys.argv[2] + batch_impl = sys.argv[2] if is_profile_batch and len(sys.argv) > 2 else "" + needs_candidate = role != "baseline" and ( + not is_profile_batch or batch_impl == "candidate") + candidate = Candidate(artifact_dir) if needs_candidate else None + + if is_profile_batch: + if batch_impl not in ("candidate", "triton"): + raise RuntimeError("profile-batch implementation must be candidate or triton") + calls = int(sys.argv[3]) if len(sys.argv) > 3 else 120 + if calls <= 0: + raise RuntimeError("profile-batch calls must be positive") + if batch_impl == "candidate" and candidate is None: + raise RuntimeError("candidate profile requested without candidate artifact") cases = _public_cases() - found = next((c for c in cases if c.id == case_id), None) - if found is None: - raise RuntimeError(f"unknown public profile case: {case_id}") - if role == "baseline": - _profile_case_triton(weights, found, device) - else: - assert candidate is not None - _profile_case(candidate, weights, found, device) + selected = { + token.strip() for token in os.environ.get( + "METAINFER_PROFILE_CASE_IDS", "" + ).split(",") if token.strip() + } + if selected: + known = {case.id for case in cases} + unknown = sorted(selected - known) + if unknown: + raise RuntimeError(f"unknown profile case ids: {unknown}") + cases = [case for case in cases if case.id in selected] + profiled_cases = [] + for found in cases: + if batch_impl == "triton": + begin_ns, end_ns, begin_epoch_ns, end_epoch_ns = _profile_batch_case_triton( + weights, found, device, calls) + else: + assert candidate is not None + begin_ns, end_ns, begin_epoch_ns, end_epoch_ns = _profile_batch_case_candidate( + candidate, weights, found, device, calls) + profiled_cases.append({ + "id": found.id, "m": found.m, "n": found.n, "k": found.k, + "host_monotonic_begin_ns": begin_ns, + "host_monotonic_end_ns": end_ns, + "host_epoch_begin_ns": begin_epoch_ns, + "host_epoch_end_ns": end_epoch_ns, + }) write_json(report_path, { "passed": True, - "case_id": found.id, - "implementation": "triton" if role == "baseline" else "candidate", - "timed_scope": "launch_w8a8_gemm_only", + "implementation": batch_impl, + "calls_per_case": calls, + "case_ids": [case.id for case in cases], + "cases": profiled_cases, + "timed_scope": "core implementation launches only", }) return @@ -592,37 +534,6 @@ def main() -> None: write_json(report_path, report) return - if phase == "benchmark": - protocol = json.loads(_env("METAINFER_BENCHMARK_PROTOCOL")) - warmup = int(protocol["warmup"]) - samples = int(protocol["samples"]) - all_cases = _public_cases() - cases_out = [] - for c in all_cases: - try: - if role == "baseline": - item = _benchmark_case_triton(weights, c, device, warmup, samples) - else: - assert candidate is not None - item = _benchmark_case_candidate( - candidate, weights, c, device, warmup, samples - ) - except Exception as exc: - write_json( - report_path, - {"passed": False, "reason": str(exc), "cases": []}, - ) - sys.exit(2) - cases_out.append(item) - write_json(report_path, { - "passed": True, - "methodology": protocol, - "timed_scope": "launch_w8a8_gemm_only", - "activation_quantization_timed": False, - "weight_loading_or_preprocessing_timed": False, - "cases": cases_out, - }) - def write_json(path: Path, data: Dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) diff --git a/metainfer/tasks/opt_GEMM_kernel/harness/user_gemm/run_hipprof_suite.py b/metainfer/tasks/opt_GEMM_kernel/harness/user_gemm/run_hipprof_suite.py new file mode 100644 index 00000000..dfb92546 --- /dev/null +++ b/metainfer/tasks/opt_GEMM_kernel/harness/user_gemm/run_hipprof_suite.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python3 +"""Task-local Custom/Triton hipprof trace and PMC collection. + +All paths come from this frozen evaluator bundle and METAINFER_* runtime +inputs. No external benchmark checkout or prebuilt kernel path is used. +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import time +from pathlib import Path + + +HERE = Path(__file__).resolve().parent +EVALUATE = HERE / "evaluate.py" + + +def _required_env(name: str) -> str: + value = os.environ.get(name) + if not value: + raise RuntimeError(f"missing environment variable: {name}") + return value + + +def _clear(base: Path) -> None: + for suffix in ("", ".db", ".csv", ".hipkernel.csv", ".hiptrace.csv"): + path = Path(str(base) + suffix) + if path.exists(): + path.unlink() + + +def _run_pass( + hipprof: Path, + output_dir: Path, + label: str, + profiler_args: list[str], + implementation: str, + calls: int, + timeout_s: int, +) -> dict[str, object]: + base = output_dir / label + _clear(base) + manifest = output_dir / f"{label}-harness.json" + env = dict(os.environ) + env.update({ + "METAINFER_EVALUATION_PHASE": "profile-batch", + "METAINFER_EVALUATION_ROLE": ( + "baseline" if implementation == "triton" else "candidate" + ), + "METAINFER_REPORT_PATH": str(manifest), + }) + command = [ + str(hipprof), *profiler_args, "-o", str(base), + sys.executable, str(EVALUATE), "profile-batch", implementation, str(calls), + ] + print("+", " ".join(command), flush=True) + started_at = time.time() + started_monotonic = time.monotonic() + try: + completed = subprocess.run( + command, cwd=HERE, env=env, text=True, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, check=False, + timeout=timeout_s, + ) + except subprocess.TimeoutExpired as exc: + output = exc.stdout or "" + if isinstance(output, bytes): + output = output.decode(errors="replace") + (output_dir / f"{label}.log").write_text(output, encoding="utf-8") + raise RuntimeError(f"{label} timed out after {timeout_s}s") from exc + (output_dir / f"{label}.log").write_text( + completed.stdout or "", encoding="utf-8") + print(completed.stdout or "", flush=True) + if completed.returncode: + raise RuntimeError(f"{label} failed with status {completed.returncode}") + report = json.loads(manifest.read_text(encoding="utf-8")) + if report.get("passed") is not True: + raise RuntimeError(f"{label} harness did not report success") + return { + "label": label, + "command": command, + "started_at": started_at, + "ended_at": time.time(), + "duration_s": time.monotonic() - started_monotonic, + "timeout_s": timeout_s, + "calls_per_case": calls, + "case_count": len(report.get("cases") or []), + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--hipprof", default="/opt/dtk/bin/hipprof") + parser.add_argument("--output-dir", type=Path) + parser.add_argument("--pmc-calls", type=int, default=1) + parser.add_argument("--passes", choices=("trace", "diagnostic", "full"), default="full") + parser.add_argument("--trace-timeout-s", type=int, default=600) + parser.add_argument("--pmc-timeout-s", type=int, default=480) + parser.add_argument("--case-ids", default="") + parser.add_argument( + "--implementations", default="candidate,triton", + help="comma-separated subset of candidate,triton") + args = parser.parse_args() + protocol = json.loads(_required_env("METAINFER_BENCHMARK_PROTOCOL")) + warmup = int(protocol["warmup"]) + samples = int(protocol["samples"]) + trace_calls = int(protocol["trace_calls"]) + if trace_calls != warmup + samples or samples < 3 or args.pmc_calls <= 0: + raise RuntimeError( + "frozen protocol requires trace_calls=warmup+samples and positive counts" + ) + + # Validate frozen runtime inputs before starting expensive profiler passes. + _required_env("METAINFER_WEIGHT_BUNDLE") + _required_env("METAINFER_BUILD_ARTIFACT_DIR") + root = args.output_dir + if root is None: + report = Path(_required_env("METAINFER_REPORT_PATH")).resolve() + root = report.parent / "hipprof-suite" + root = root.resolve() + root.mkdir(parents=True, exist_ok=True) + hipprof = Path(args.hipprof).resolve() + if not hipprof.is_file(): + raise RuntimeError(f"hipprof not found: {hipprof}") + + implementations = tuple( + token.strip() for token in args.implementations.split(",") + if token.strip()) + if not implementations or any( + value not in ("candidate", "triton") for value in implementations + ): + raise RuntimeError("implementations must contain candidate and/or triton") + + case_ids = [token.strip() for token in args.case_ids.split(",") if token.strip()] + if case_ids: + os.environ["METAINFER_PROFILE_CASE_IDS"] = ",".join(case_ids) + pass_records: list[dict[str, object]] = [] + + for impl in implementations: + if args.passes in ("trace", "full"): + pass_records.append(_run_pass( + hipprof, root, f"{impl}-trace", ["--hip-trace", "--stats"], + impl, trace_calls, args.trace_timeout_s)) + if args.passes in ("diagnostic", "full"): + for label, mode in ( + ("pmc", "--pmc"), + ("read", "--pmc-read"), + ("write", "--pmc-write"), + ): + pass_records.append(_run_pass( + hipprof, root, f"{impl}-{label}", + [mode, "--pmc-type", "3"], impl, args.pmc_calls, + args.pmc_timeout_s)) + + (root / "collection.json").write_text(json.dumps({ + "passed": True, + "implementations": implementations, + "trace_calls": trace_calls, + "warmup": warmup, + "samples": samples, + "trace_timing": "arithmetic mean after frozen warmup calls", + "pmc_calls": args.pmc_calls, + "passes": args.passes, + "case_ids": case_ids, + "pass_records": pass_records, + "weight_bundle": str(Path(_required_env("METAINFER_WEIGHT_BUNDLE")).resolve()), + "artifact_dir": str(Path(_required_env("METAINFER_BUILD_ARTIFACT_DIR")).resolve()), + }, indent=2) + "\n", encoding="utf-8") + print(f"wrote profiler suite to {root}") + + +if __name__ == "__main__": + main() diff --git a/metainfer/tasks/opt_GEMM_kernel/harness/user_gemm/task.yaml b/metainfer/tasks/opt_GEMM_kernel/harness/user_gemm/task.yaml index e8d29537..24d09df8 100644 --- a/metainfer/tasks/opt_GEMM_kernel/harness/user_gemm/task.yaml +++ b/metainfer/tasks/opt_GEMM_kernel/harness/user_gemm/task.yaml @@ -48,9 +48,6 @@ commands: correctness: argv: ["python3", "{bundle_dir}/evaluate.py", correctness] timeout_s: 7200 - benchmark: - argv: ["python3", "{bundle_dir}/evaluate.py", benchmark] - timeout_s: 7200 profile: argv: ["python3", "{bundle_dir}/evaluate.py"] timeout_s: 1800 @@ -70,12 +67,6 @@ cases: benchmark: matrix: m_values: [1, 2, 4, 8, 16, 4096] - large_m: 4096 - small_m_total_weight: 0.5 - large_m_weight: 0.5 - # Every requested M is protected from >3% regression. This prevents a - # large-M win from hiding a decode regression (or vice versa). - critical_m: [1, 2, 4, 8, 16, 4096] workloads: - {id: wqkv-a-tp4, op: wqkv_a, tp: 4, k: 4096, n: 1536} - {id: wq-b-tp4, op: wq_b, tp: 4, k: 1024, n: 8192} @@ -91,13 +82,18 @@ cases: benchmark_protocol: warmup: 10 samples: 100 - timer: gpu_event - statistic: median - synchronization: event_per_sample - timed_scope: launch_w8a8_gemm_only + trace_calls: 110 + timer: hipprof_gpu_kernel_duration_ns + statistic: arithmetic_mean + operator_aggregation: sum_gpu_kernel_duration_per_call + synchronization: hipprof_trace + timed_scope: operator_gpu_dispatches_only + host_launch_time_included: false + pmc_timing_used: false + raw_samples_retained: true + dispersion_reported: [median, stddev, cv, min, max] + comparison: same_round_champion + boundary_retest: equal_sample_second_trace_for_both_sides acceptance: - min_weighted_speedup: 1.01 noise_threshold: 0.01 - max_critical_regression: 0.03 - require_all_cases: true diff --git a/metainfer/tasks/opt_GEMM_kernel/notebooks/00_task_contract.md b/metainfer/tasks/opt_GEMM_kernel/notebooks/00_task_contract.md index ee88d75d..857c49ee 100644 --- a/metainfer/tasks/opt_GEMM_kernel/notebooks/00_task_contract.md +++ b/metainfer/tasks/opt_GEMM_kernel/notebooks/00_task_contract.md @@ -1,30 +1,35 @@ # GEMM optimization contract -The candidate implements the GEMM family described in the task requirements: +The candidate implements the GEMM family defined by the frozen evaluator: ```text C = epilogue(alpha * op(A) @ op(B) + beta * C_or_bias) ``` -The exact dtype, transpose flags, layouts, strides, batching, alignment, -epilogue, legal approximation and workspace limits come from the user task and -the evaluator bundle. `task.yaml::public_contract` is the frozen source of -truth supplied to agents and shown read-only in the UI. Unspecified behavior -must not be guessed silently. +Exact dtype, transpose flags, layouts, strides, batching, alignment, numerics, +epilogue, legal approximation, and workspace limits come from +`task.yaml::public_contract`. It is the source of truth supplied to agents and +shown read-only in the UI. Unspecified behavior must not be guessed silently. Acceptance requires all of the following: -- the system compiler command succeeds; -- every declared public and held-out correctness case is returned and passes; -- every performance case is returned under one fixed timing methodology; -- trace-weighted speedup clears the configured minimum; -- no critical shape exceeds its regression limit; -- the candidate beats the current champion by more than the configured noise - threshold. +- the system-owned build succeeds; +- every public and held-out correctness case is returned and passes; +- hipprof returns one finite positive operator latency for every performance + shape under the exact frozen methodology; +- each shape is strictly faster than the frozen Triton baseline; +- each shape crosses the current Champion by the configured noise threshold. -Before the optimization loop starts, the original submission must compile, -pass every correctness case, and produce a complete benchmark under the frozen -BuildProfile. This certified measurement is the only baseline used later. +Operator latency is the arithmetic mean of the final trace samples after summing +all GPU dispatch `DurationNs` belonging to each logical GEMM call. Host launch, +JIT, allocation, copies, preprocessing, and synchronization are outside timing. +PMC replay supplies diagnostics only and never supplies latency. -Only files under `submission/` are candidate deliverables. Agent-written test -or benchmark scripts are useful local diagnostics but never become gates. +There are no performance weights, critical-shape exceptions, or aggregate score +that can compensate for a losing shape. Before optimization, Triton and Initial +HIP are independently built, correctness-checked, and profiled under the frozen +BuildProfile. Their immutable report references are the only performance facts +used later. + +Only files under `submission/` are candidate deliverables. Agent-written tests, +benchmarks, profiler commands, or pass/fail logic never become system gates. diff --git a/metainfer/tasks/opt_GEMM_kernel/notebooks/02_evaluation_protocol.md b/metainfer/tasks/opt_GEMM_kernel/notebooks/02_evaluation_protocol.md index 61158c3d..b8d470a7 100644 --- a/metainfer/tasks/opt_GEMM_kernel/notebooks/02_evaluation_protocol.md +++ b/metainfer/tasks/opt_GEMM_kernel/notebooks/02_evaluation_protocol.md @@ -1,111 +1,150 @@ # Fixed evaluation protocol -The task snapshots an external evaluator bundle into task state before the -first agent runs. A SHA-256 manifest is checked before and after every system -evaluation command. Compilation is not an evaluator command: MetaInfer owns -the frozen BuildProfile, CMakeLists.txt and build.sh. The evaluator bundle -contains only correctness and benchmark commands: +The task snapshots the evaluator and weight bundles before any agent runs. Their +SHA-256 manifests are verified at system gates. MetaInfer separately owns the +BuildProfile, generated CMake, compiler, GPU architecture, hipprof command, and +counter groups. + +The evaluator `task.yaml` contains the public contract, correctness command, +profile entry point, correctness cases, exact benchmark shapes, and frozen +hipprof protocol. It does not assign performance weights or criticality: ```yaml schema_version: 2 name: example-gemm public_contract: - operation: "C = alpha * A @ B + beta * C" - dtype: {a: fp16, b: fp16, accumulation: fp32, c: fp16} - layout: - {a: row_major, b: row_major, c: row_major, trans_a: false, trans_b: false} + operation: "Y = scaled_int8_gemm(A, W, A_scale, W_scale)" + dtype: {a: int8, b: int8, accumulation: int32, c: bfloat16} + layout: {a: row_major, b: row_major, c: row_major} abi: - entrypoint: launch_gemm - signature: "launch_gemm(A, B, C, M, N, K, stream)" + entrypoint: launch_w8a8_gemm + signature: "launch_w8a8_gemm(A, W, A_scale, W_scale, Y, M, N, K, stream)" commands: correctness: argv: [python3, evaluate.py, correctness] - timeout_s: 1200 - benchmark: - argv: [python3, evaluate.py, benchmark] + timeout_s: 7200 + profile: + argv: [python3, evaluate.py] timeout_s: 1800 cases: - correctness: [public-1, public-2, heldout-1] + correctness: [public-1, heldout-1] private: [heldout-1] benchmark: - id: decode-gemm - weight: 2000 - critical: true - shape: {m: 1, n: 4096, k: 4096, batch: 1} - bytes: 33570816 + shape: {m: 1, n: 4096, k: 4096} + bytes: 16797696 - id: prefill-gemm - weight: 100 - critical: false - shape: {m: 2048, n: 4096, k: 4096, batch: 1} - bytes: 67108864 + shape: {m: 4096, n: 4096, k: 4096} + bytes: 50331648 benchmark_protocol: warmup: 10 samples: 100 - timer: gpu_event + trace_calls: 110 + timer: hipprof_gpu_kernel_duration_ns + statistic: arithmetic_mean + operator_aggregation: sum_gpu_kernel_duration_per_call + synchronization: hipprof_trace + timed_scope: operator_gpu_dispatches_only + host_launch_time_included: false + pmc_timing_used: false acceptance: - min_weighted_speedup: 1.01 noise_threshold: 0.01 - max_critical_regression: 0.03 - require_all_cases: true ``` -Each command writes JSON to `METAINFER_REPORT_PATH`. +`public_contract` is the only source of truth for dtype, layout, numerics, and +candidate ABI. Benchmark `shape` is mandatory. Frozen optional `flops` and +`bytes` metadata is used only to derive diagnostic TFLOPS or modeled bandwidth; +it does not affect pass/fail. -`public_contract` is mandatory and is the only source of truth for dtype, -layout and candidate ABI. Benchmark case `shape` is mandatory. The creation UI -does not ask the task owner to duplicate these fields: after the evaluator is -frozen, the task detail page renders the extracted contract read-only and the -orchestrator injects exactly the same contract into planner/implementer -prompts. +## Correctness and performance reports -Correctness report: +The correctness command writes JSON to `METAINFER_REPORT_PATH`: ```json { "passed": true, "cases": [ - {"id": "public-1", "passed": true, "max_abs_error": 0.001} + {"id": "public-1", "passed": true, "max_abs_error": 0.0} ] } ``` -Benchmark report: +Performance is not supplied by an agent-authored benchmark command. The +system-owned profiler runs the frozen task-local hipprof suite and constructs a +canonical benchmark report from trace operator times: ```json { + "schema_version": 2, "passed": true, - "methodology": {"warmup": 10, "samples": 100, "timer": "gpu_event"}, + "methodology": { + "warmup": 10, + "samples": 100, + "trace_calls": 110, + "timer": "hipprof_gpu_kernel_duration_ns", + "statistic": "arithmetic_mean", + "operator_aggregation": "sum_gpu_kernel_duration_per_call", + "synchronization": "hipprof_trace", + "timed_scope": "operator_gpu_dispatches_only", + "host_launch_time_included": false, + "pmc_timing_used": false + }, + "timing_source": "hipprof GPU kernel DurationNs", + "timed_scope": "operator_gpu_dispatches_only", + "profile_report": { + "path": "logs/001/candidate-hardware-profile.json", + "sha256": "..." + }, "cases": [ { "id": "decode-gemm", - "latency_ms": 0.11 + "latency_ms": 0.011, + "shape": {"m": 1, "n": 4096, "k": 4096}, + "dispatch_count": 2, + "kernel_breakdown_us": {"split": 8.0, "reduce": 3.0} } ] } ``` -Before any optimizer runs, the system compiles the original submission and -runs correctness and benchmark with `METAINFER_EVALUATION_ROLE=baseline`. -That report and its BuildProfile fingerprint are frozen. Candidate runs use -`role=candidate`; they only report their own latency. Weight and criticality -come from task.yaml, not from measurement reports. +The methodology must exactly match the frozen protocol. Expected case IDs must +have a one-to-one mapping to finite positive latency values; missing, duplicate, +or unexpected cases fail validation. Each logical call's latency is the sum of +all related GPU dispatch `DurationNs`. PMC replay duration cannot populate +`latency_ms`. -For GEMM profiler display, each benchmark case may declare `shape` and -`bytes`. When `shape` is present, the frozen spec derives FLOPs as -`2 * M * N * K * batch`; an explicit positive `flops` value overrides that -derivation for fused or non-standard work. `bytes` is the task author's -declared total device-memory traffic for the case and should include every -tensor read/write required by the ABI. Candidate reports never provide these -values. +## Every-shape gates -The methodology object must exactly match `benchmark_protocol` for both -baseline and candidate. The orchestrator computes weighted speedup as: +For every frozen benchmark case: ```text -sum(weight_i * baseline_ms_i) / sum(weight_i * candidate_ms_i) +candidate_ms < triton_baseline_ms +candidate_ms < champion_ms * (1 - noise_threshold) ``` -The UI derives profiler rates from the frozen work metadata and measured +Champion evaluation uses a strict boundary where required by promotion so an +equality at the threshold cannot become a hidden improvement. A failure on any +shape rejects the candidate. `worst_case_speedup` and failed IDs are diagnostics, +not aggregate substitutes for the gate. + +## Performance report as source of truth + +Canonical reports are written atomically and referenced by task-state-relative +path plus SHA-256: + +```text +baseline/baseline-benchmark-report.json +certified/initial-hip/candidate-benchmark-report.json +logs//candidate-benchmark-report.json +``` + +The Champion v2 record stores its submission digest and measurement-report +reference, not copied per-shape latency or an aggregate score. Promotion and +cold restart verify the digest and reload the referenced report. Iteration score, +timeline, and API summaries are derived historical views and cannot drive a +future promotion. + +The UI derives optional rates only from frozen work metadata and authoritative latency: ```text @@ -113,5 +152,5 @@ TFLOPS = flops / latency_ms / 1e9 GB/s = bytes / latency_ms / 1e6 ``` -If `shape`/`flops` or `bytes` is omitted, latency and speedup remain valid and -the corresponding TFLOPS or bandwidth tile is shown as unavailable. +If optional work metadata is absent, latency and the all-shape gate remain valid +while the corresponding rate is unavailable. diff --git a/metainfer/tasks/opt_GEMM_kernel/notebooks/04_profiling.md b/metainfer/tasks/opt_GEMM_kernel/notebooks/04_profiling.md index d22794ba..6e41c18b 100644 --- a/metainfer/tasks/opt_GEMM_kernel/notebooks/04_profiling.md +++ b/metainfer/tasks/opt_GEMM_kernel/notebooks/04_profiling.md @@ -1,10 +1,9 @@ # K100 / gfx928 fixed profiling route -The WebUI's `Hygon K100` selection resolves one system-owned execution -profile. Agents do not select tools or construct commands. +The Web UI's `Hygon K100` selection resolves one system-owned build and profile. +Agents do not select tools, construct commands, or provide timing results. -The compilation route is equivalent to the C++ framework task's hardware -binding: +## Frozen build ```text cmake -S system_build -B ITER_BUILD -G Ninja \ @@ -13,68 +12,89 @@ cmake --build ITER_BUILD --target \ metainfer_gemm_candidate metainfer_gemm_harness ``` -The generated CMake freezes the resolved DTK `hipcc`, Release `-O3`, C++/HIP -17, and `HIP_ARCHITECTURES=gfx928`. Exact resolved paths, versions, commands, -flags and the profile fingerprint are written to the compile report. +The generated build freezes the resolved DTK `hipcc`, Release `-O3`, C++/HIP 17, +and `HIP_ARCHITECTURES=gfx928`. Resolved paths, versions, flags, architecture, +and the BuildProfile fingerprint are written to the compile report. -E first consumes the Harness GPU-event benchmark for every weighted shape. -It then invokes the Harness as `profile CASE_ID` for M=1, M=16 and M=4096 of -the public `wq_b TP=4` workload. On the K100 DTK installation, the preferred -system command is: +## Required hipprof suite -```text -hipprof --pmc --pmc-type 3 -o \ - metainfer_gemm_harness profile -``` - -`--pmc-type 3` produces a CSV table. Per-instance columns such as -`TCC_HIT[0..31]` and `TCC_MISS[0..31]` are summed by MetaInfer before L2 rates -are derived. Multiple dispatches (for example split-K plus its reduction) are -retained in the normalized case report. PMC timings are diagnostic only and do -not replace the Harness GPU-event benchmark. A PMC CSV is accepted only when -the same invocation writes a successful `harness-profile.json` whose -`case_id` exactly matches the requested case. Since the profile entrypoint -performs all preparation and synchronization before its single candidate -launch, the captured kernel dispatches belong to that case; a split-K main -kernel and its reduction are deliberately retained together. - -If hipprof is unavailable, rocprofv3 has this fixed shape: +K100 performance uses only the frozen task-local hipprof suite: ```text -rocprofv3 --pmc --output-format csv json \ - --output-directory \ - --kernel-include-regex w8a8_scaled_ -- \ - metainfer_gemm_harness profile + /run_hipprof_suite.py \ + --hipprof /opt/dtk/bin/hipprof --output-dir + /analyze_hipprof_suite.py ``` -For a DTK installation that provides legacy rocprof, the fixed fallback is: +The active K100 profile accepts hipprof only. A missing executable, suite, +analyzer, shape, pass, or matching profiler/protocol fingerprint is an +infrastructure failure. There is no GPU Event, rocprofv3, legacy rocprof, or +PMC-duration latency fallback. -```text -rocprof -i -o \ - --timestamp on metainfer_gemm_harness profile -``` +For both Triton and candidate, the suite performs one trace collection and +separate `--pmc`, `--pmc-read`, and `--pmc-write` collections. Tensor generation, +quantization, weight loading/packing, workspace allocation, candidate setup, +Triton JIT, and synchronization complete before each marked host interval. Only +repeated steady-state GEMM calls occur inside the interval. + +## Trace operator latency -The available-counter query is performed once when the profile is frozen; -unsupported names are removed from the whitelist rather than guessed. Tool -path, version, counter groups and representative shapes are fingerprinted. -Profiler failure is an E-stage infrastructure failure for this K100 profile. +Every shape has 110 trace calls. The first 10 are warmup and the final 100 are +measured. hipprof trace rows are selected using the manifest's realtime host +boundaries. Legacy manifests without realtime boundaries may translate their +monotonic interval with a boot-stable realtime-minus-monotonic offset; a warmup +kernel timestamp is not used to infer that offset. -The hardware profile is diagnostic evidence for F. Champion promotion remains -owned by correctness plus the complete weighted multi-shape event benchmark; -the three profiler cases do not replace or reweight that score. +The analyzer verifies: -# Interpretation checklist +1. trace, PMC, read, and write manifests contain the same case IDs and M/N/K; +2. manifest call counts match the frozen collection protocol; +3. selected dispatch counts divide exactly into logical calls; +4. final measured calls have a stable kernel dispatch pattern. -Record the target GPU and exact compiler flags before interpreting profiler -data. Useful signals include achieved occupancy, waves/SM or waves/CU, register -and shared-memory pressure, memory transaction efficiency, cache hit rate, -tensor-core/MFMA utilization, synchronization stalls and launch count. +For one logical call: -Do not optimize a single profiler counter in isolation. A lower occupancy -kernel may still win through better instruction-level parallelism or data -reuse. Conversely, a headline speedup smaller than run-to-run noise is not a -promotion. +```text +operator_us = sum(DurationNs of every related GPU dispatch) / 1000 +``` -Use public per-shape results to identify the class that changed. Held-out -results are deliberately summarized so implementation choices generalize -rather than overfit case IDs. +The reported latency is the arithmetic mean of the final 100 `operator_us` +values. Repeated same-name dispatches are first summed within a call, then their +per-call contributions are averaged for the kernel breakdown. The longest +kernel name is only a resource-label hint; it never replaces operator latency. +Host launch API and synchronization time are outside this metric. + +## PMC diagnostics + +DTK hipprof `--pmc-type 3` internally replays hardware counter groups and emits +merged indexed columns per original dispatch. The analyzer sums indexed values +such as `TCC_HIT[0..N]` and `TCC_MISS[0..N]`, and aggregates every operator +dispatch before normalizing by logical call count. + +Separate read/write passes derive physical HBM request bytes and bandwidth. +The compact report retains, when actually reported: + +- HBM read/write bytes and read/write/total GB/s; +- L2 hit percentage; +- VGPR, AGPR, SGPR, LDS, and scratch for the selected resource-label kernel; +- grid size, workgroup size, wave size, waves per workgroup, and dispatch count; +- occupancy or wave-residency only when hipprof exposes a reliable field. + +Replay `DurationNs` or `DispatchNs` is instrumentation time and is never copied +into benchmark latency. `occupancy_pct` remains unavailable rather than being +estimated from incomplete metadata. + +## Interpretation checklist + +Start from every shape's summed operator latency and dispatch breakdown. Use PMC +to test a bounded hypothesis, for example excessive physical HBM traffic, weak +L2 reuse, high register/LDS pressure, partial/reduction overhead, or insufficient +parallelism. Do not optimize one counter in isolation: lower occupancy can win +through data reuse or instruction-level parallelism, while higher bandwidth can +still lose if it increases dispatch or reduction work. + +A performance improvement is accepted only if every frozen shape beats Triton +and every shape crosses the current Champion noise threshold. Representative +cases may guide diagnosis, but the profiler report and promotion gate retain all +60 shapes. Notebook timings and older profiler records are historical evidence, +not current service-level targets. diff --git a/metainfer/tasks/opt_GEMM_kernel/notebooks/05_champion_policy.md b/metainfer/tasks/opt_GEMM_kernel/notebooks/05_champion_policy.md index 89b9282d..73bd1dc0 100644 --- a/metainfer/tasks/opt_GEMM_kernel/notebooks/05_champion_policy.md +++ b/metainfer/tasks/opt_GEMM_kernel/notebooks/05_champion_policy.md @@ -1,14 +1,29 @@ # Champion/challenger policy -Every iteration starts from the persisted champion, not merely the most recent -candidate. A challenger is promoted only after compile, complete correctness, -multi-shape scoring and critical-regression gates pass. +Every HIP iteration starts from the persisted HIP Champion. While Triton remains +Champion, iterations start from certified Initial HIP because Triton has no +editable HIP submission tree. -The challenger must also exceed the champion's weighted speedup by the noise -threshold. Failed and non-promoted candidates remain in iteration history for -diagnosis, but they never become the starting implementation for the next -iteration. +A challenger is eligible only after the system build and every correctness case +pass. Its immutable hipprof performance report must contain exactly one finite +positive operator latency for every frozen benchmark shape. Promotion then +requires both per-shape gates: -At the end of the task, `state/champion/submission/` is the selected artifact -and `champion.json` identifies its source iteration and score. +```text +candidate_ms < triton_baseline_ms +candidate_ms < champion_ms * (1 - noise_threshold) +``` +The strict promotion comparison rejects equality where it would not represent a +real improvement. One failed shape rejects the challenger; no weighted mean, +critical-shape exception, or favorable aggregate can compensate. + +`champion.json` v2 stores Champion kind, source iteration, submission SHA-256, +promotion metadata, and a task-state-relative measurement-report path plus +SHA-256. It does not copy per-shape latency or aggregate score. Promotion and +cold restart verify and reload that report. Iteration score and timeline values +are derived historical snapshots only. + +Failed and non-promoted candidates remain in iteration history for diagnosis but +never become the next starting implementation. The selected HIP artifact is +stored under `champion/submission/`; a Triton Champion has no copied HIP source. diff --git a/metainfer/tasks/opt_GEMM_kernel/notebooks/09_small_M_splitK_sdot4.md b/metainfer/tasks/opt_GEMM_kernel/notebooks/09_small_M_splitK_sdot4.md index 923863f9..3f9f65fe 100644 --- a/metainfer/tasks/opt_GEMM_kernel/notebooks/09_small_M_splitK_sdot4.md +++ b/metainfer/tasks/opt_GEMM_kernel/notebooks/09_small_M_splitK_sdot4.md @@ -1,5 +1,11 @@ # K100/gfx928 小 M、大 K：128-bit Load + Split-K + SDOT4 +> **历史方法说明**：本文保留当时 GPU Event/旧 profiler 的实验数字与推导，便于追踪 +> 技术来源，但这些计时不再是当前任务的 latency 或 promotion 证据。当前协议只使用 +> frozen hipprof trace：每次逻辑调用汇总全部 GPU dispatch `DurationNs`，对最终样本取 +> 算术平均；PMC replay 仅作诊断；每个 frozen shape 都必须独立通过。本文的 dispatch +> 建议只能作为待复验假设。 + ## 1. 结论 目标算子： diff --git a/metainfer/tasks/opt_GEMM_kernel/notebooks/10_gfx928_MMAC_tensorcore_general_GEMM.md b/metainfer/tasks/opt_GEMM_kernel/notebooks/10_gfx928_MMAC_tensorcore_general_GEMM.md index 3b269be6..e9a1df8c 100644 --- a/metainfer/tasks/opt_GEMM_kernel/notebooks/10_gfx928_MMAC_tensorcore_general_GEMM.md +++ b/metainfer/tasks/opt_GEMM_kernel/notebooks/10_gfx928_MMAC_tensorcore_general_GEMM.md @@ -1,5 +1,11 @@ # K100/gfx928 INT8 MMAC（TensorCore）通用 GEMM 与 Split-K 选择 +> **历史方法说明**：本文保留当时 GPU Event/旧 profiler 的实验数字与推导，便于追踪 +> 技术来源，但这些计时不再是当前任务的 latency 或 promotion 证据。当前协议只使用 +> frozen hipprof trace：每次逻辑调用汇总全部 GPU dispatch `DurationNs`，对最终样本取 +> 算术平均；PMC replay 仅作诊断；每个 frozen shape 都必须独立通过。本文的 dispatch +> 建议只能作为待复验假设。 + ## 1. 结论 目标算子： diff --git a/metainfer/tasks/opt_GEMM_kernel/notebooks/11_champion_engineering_DPP_alignment_generality.md b/metainfer/tasks/opt_GEMM_kernel/notebooks/11_champion_engineering_DPP_alignment_generality.md index 9be09536..af72ed63 100644 --- a/metainfer/tasks/opt_GEMM_kernel/notebooks/11_champion_engineering_DPP_alignment_generality.md +++ b/metainfer/tasks/opt_GEMM_kernel/notebooks/11_champion_engineering_DPP_alignment_generality.md @@ -1,5 +1,11 @@ # K100/gfx928 Champion 工程化：DPP、Split-K、MMAC、安全对齐与通用性 +> **历史方法说明**：本文保留当时 GPU Event/旧 profiler 的实验数字与推导，便于追踪 +> 技术来源，但这些计时不再是当前任务的 latency 或 promotion 证据。当前协议只使用 +> frozen hipprof trace：每次逻辑调用汇总全部 GPU dispatch `DurationNs`，对最终样本取 +> 算术平均；PMC replay 仅作诊断；每个 frozen shape 都必须独立通过。本文的 dispatch +> 建议只能作为待复验假设。 + ## 1. 文档目的与源码基线 本文记录对以下算子的实际修改、失败尝试、修复、性能证据和工程边界，供后续 planning、implementer 和 reviewer agent 直接参考： diff --git a/metainfer/tasks/opt_GEMM_kernel/notebooks/12_MMAC_CTA_swizzle_fused_splitK_BM1.md b/metainfer/tasks/opt_GEMM_kernel/notebooks/12_MMAC_CTA_swizzle_fused_splitK_BM1.md index eef7f6c7..405862eb 100644 --- a/metainfer/tasks/opt_GEMM_kernel/notebooks/12_MMAC_CTA_swizzle_fused_splitK_BM1.md +++ b/metainfer/tasks/opt_GEMM_kernel/notebooks/12_MMAC_CTA_swizzle_fused_splitK_BM1.md @@ -1,5 +1,11 @@ # K100/gfx928 W8A8 GEMM：MMAC CTA Swizzle、Fused Split-K 与 BM=1 特化 +> **历史方法说明**：本文保留当时 GPU Event/rocprof 的实验数字与推导，便于追踪技术 +> 来源，但这些计时不再是当前任务的 latency 或 promotion 证据。当前协议只使用 frozen +> hipprof trace：每次逻辑调用汇总全部 GPU dispatch `DurationNs`，对最终样本取算术 +> 平均；PMC replay 仅作诊断；每个 frozen shape 都必须独立通过。本文的 dispatch 建议 +> 只能作为待复验假设。 + ## 1. 文档范围与证据基线 本文记录在 K500SM_AI / gfx928 / Wave64 上实测过的三项改动： diff --git a/metainfer/tasks/opt_GEMM_kernel/notebooks/13_stream_splitK workload_opt.md b/metainfer/tasks/opt_GEMM_kernel/notebooks/13_stream_splitK workload_opt.md new file mode 100644 index 00000000..09f94983 --- /dev/null +++ b/metainfer/tasks/opt_GEMM_kernel/notebooks/13_stream_splitK workload_opt.md @@ -0,0 +1,548 @@ +# Split-K 按 Stream 隔离 Workspace：实现说明 + +> **历史适用范围与验证状态**：本文记录的是 `001_swizzle.cpp` 的 +> per-stream 内部 workspace/fused Split-K 方案，不是 `/home/FF/workspace/003` +> 最终采用的 caller-owned persistent workspace ABI，也不是当前任务必须遵循的 +> dispatch recipe。文末列出的 gfx928 编译、双 stream/多 device 正确性压力测试和 +> 新旧交替 hipprof benchmark 尚未完成。Agent 只能把这里的并发风险、ticket +> 不变量和候选实现当作证据；必须先检查当前 submission、shape、真实 hipprof +> operator time/dispatch breakdown 和 PMC，再决定是否移植、修改或完全舍弃此方案。 + +本文档记录 `001_swizzle.cpp` 中新增的 Split-K host 侧改造，供后续 +agent 继续实现、移植和验证。 + +对应源码： + +- `001_swizzle.cpp` +- fused Split-K kernel：`small_m_splitk_dot4_fused_kernel` +- host 模板启动器：`launch_splitk_fused_instance` + +## 1. 改造目标 + +原实现只有一组进程级静态设备指针： + +```cpp +static int32_t* g_workspace; +static uint32_t* g_tile_done; +``` + +所有 Split-K 调用都复用这两个地址。不同 HIP stream 上的 kernel 可以 +并发执行，因此两个请求可能同时覆盖相同的 partial，并把各自的 +`atomicAdd` ticket 混在一起。 + +可能结果： + +- Reduce 读到两个请求混合的 partial； +- 某个 CTA 被错误地判定为最后一个 split； +- counter 提前清零； +- 后续调用继承非零 counter； +- 输出发生偶发、非确定性错误。 + +本次改造的目标是： + +1. 每个 `(HIP device, hipStream_t)` 使用独立的 partial 和 tile counter； +2. 同一 stream 上的 host 提交线程安全； +3. 不同 stream 的 GPU kernel 仍能并发； +4. 稳态调用不执行 stream/device synchronize； +5. 不改变 BM、BN、BK、split 数和 device 计算逻辑。 + +## 2. 内存所有权 + +算子只为 Split-K 中间结果分配设备内存： + +```text +partial: + split_k * M * N * sizeof(int32_t) + +tile_done: + ceil(M / BM) * ceil(N / BN) * sizeof(uint32_t) +``` + +以下内存由调用方管理，算子不为其执行 `hipMalloc`： + +```text +x_q A 矩阵 +weight_kn B 矩阵 +x_scale +weight_scale +output_bf16 +``` + +kernel 内的 `a_tile`、`b_tile` 是每个 CTA 自动分配的 LDS；accumulator +属于寄存器，也不在 host workspace 中。 + +由于 kernel launch 是异步的，调用方必须保证 A、B、scale 和 output +在对应 stream 完成之前保持有效。 + +## 3. Per-stream 状态 + +每个 stream 对应一个 `StreamWorkspace`： + +```cpp +struct StreamWorkspace { + int32_t* partial = nullptr; + size_t partial_capacity = 0; + + uint32_t* tile_done = nullptr; + size_t tile_done_capacity = 0; + + std::mutex launch_mutex; +}; +``` + +含义： + +- `partial`：该 stream 的 Split-K INT32 partial； +- `tile_done`：该 stream 每个输出 tile 的完成计数； +- `partial_capacity`：按字节记录； +- `tile_done_capacity`：按 counter 元素数记录； +- `launch_mutex`：保护同一 stream 的扩容与 kernel 提交顺序。 + +进程级容器使用 `(device_id, stream_handle)` 找到状态。最新版不能只用 +stream 数值作为 key，因为默认 stream 在不同 device 上通常都表现为 +空句柄，同一个进程管理多个 GPU 时会发生冲突。 + +key 和 hash 的源码如下： + +```cpp +struct DeviceStreamKey { + int device; + uintptr_t stream; + + bool operator==(const DeviceStreamKey& other) const noexcept { + return device == other.device && stream == other.stream; + } +}; + +struct DeviceStreamKeyHash { + size_t operator()(const DeviceStreamKey& key) const noexcept { + const size_t device_hash = std::hash{}(key.device); + const size_t stream_hash = + std::hash{}(key.stream); + return device_hash ^ + (stream_hash + size_t{0x9e3779b9} + + (device_hash << 6) + (device_hash >> 2)); + } +}; + +static std::mutex g_stream_workspaces_mutex; + +static std::unordered_map< + DeviceStreamKey, + std::unique_ptr, + DeviceStreamKeyHash +> g_stream_workspaces; +``` + +全局 map mutex 只在查找或首次创建状态时短暂持有，不会持有到 GPU +kernel 完成。 + +## 4. 获取 StreamWorkspace + +获取 workspace 前先查询当前 HIP device，然后和原始 stream 数值共同 +组成 key： + +```cpp +__host__ static StreamWorkspace* +get_stream_workspace(hipStream_t stream) { + int device = -1; + if (hipGetDevice(&device) != hipSuccess) + return nullptr; + + const DeviceStreamKey key{ + device, + reinterpret_cast(stream) + }; + + std::lock_guard lock( + g_stream_workspaces_mutex); + + auto it = g_stream_workspaces.find(key); + if (it != g_stream_workspaces.end()) + return it->second.get(); + + auto workspace = std::make_unique(); + StreamWorkspace* result = workspace.get(); + g_stream_workspaces.emplace( + key, std::move(workspace)); + return result; +} +``` + +获取流程： + +1. 调用 `hipGetDevice` 获取当前 device； +2. 构造 `(device_id, stream_handle)`； +3. 锁住 `g_stream_workspaces_mutex`； +4. 查找 workspace； +5. 不存在时创建 `StreamWorkspace`； +6. 返回稳定的 `StreamWorkspace*`； +7. 释放全局 map mutex。 + +如果 `hipGetDevice` 失败，`get_stream_workspace` 返回空指针。模板启动器 +必须在解引用前检查： + +```cpp +StreamWorkspace* workspace = + get_stream_workspace(stream); + +if (!workspace) + return static_cast(hipErrorInvalidDevice); +``` + +map 的 value 使用 `std::unique_ptr`，因此 unordered_map rehash 后， +`StreamWorkspace` 本体地址仍保持稳定。 + +## 5. Workspace 扩容 + +`ensure_stream_workspace` 同时检查 partial 和 counter： + +```cpp +grow_partial = + required_partial_bytes > partial_capacity; + +grow_tile_done = + required_tile_count > tile_done_capacity; +``` + +容量足够时直接返回，不执行： + +- `hipStreamSynchronize`； +- `hipMalloc`； +- `hipFree`； +- `hipMemsetAsync`。 + +需要扩容时，旧指针可能仍被该 stream 中较早提交的 kernel 使用。 +因此替换旧内存前执行： + +```cpp +hipStreamSynchronize(stream); +``` + +这里只同步当前 stream，不调用 `hipDeviceSynchronize`。 + +扩容后： + +- partial 不需要初始化，因为每个有效输出元素都会被当前 split CTA + 覆盖； +- 新 counter 必须在同一 stream 中清零： + +```cpp +hipMemsetAsync( + workspace.tile_done, + 0, + tile_count * sizeof(uint32_t), + stream); +``` + +同一 stream 的后续 kernel launch 排在 memset 后面，因此首次使用时 +counter 一定为零。 + +注意：传统 `hipMalloc/hipFree` 本身可能包含运行时级同步成本。该成本 +只应出现在首次分配或容量增长阶段，不能出现在稳态热路径。 + +## 6. Host 模板启动器 + +重复的 grid、counter、workspace 和 kernel launch 逻辑被封装为： + +```cpp +template +__host__ __forceinline__ int +launch_splitk_fused_instance(...); +``` + +模板参数用于编译期确定： + +- workgroup 大小； +- LDS 数组尺寸； +- load 循环边界； +- SDOT4 循环边界； +- `__launch_bounds__`； +- 具体 device kernel 符号。 + +启动流程必须保持以下顺序： + +```text +计算 m_tiles / n_tiles / tile_count / partial_bytes + ↓ +根据 hipStream_t 获取 StreamWorkspace + ↓ +锁住 workspace.launch_mutex + ↓ +ensure_stream_workspace + ↓ +在同一 stream 启动 fused Split-K kernel + ↓ +hipGetLastError + ↓ +释放 workspace.launch_mutex +``` + +`launch_mutex` 只保护 host 侧的资源变更和 enqueue，不等待 kernel +执行结束。 + +同一 stream 的两个 kernel 依靠 HIP stream FIFO 自动串行；不同 +stream 使用不同设备指针，可以在 GPU 上重叠执行。 + +## 7. 编译期 Launch Bounds + +Split-K kernel 从固定： + +```cpp +__launch_bounds__(512) +``` + +改为： + +```cpp +template +__global__ __launch_bounds__(BM * BN) +void small_m_splitk_dot4_fused_kernel(...); +``` + +当前实例： + +```text +BM=1, BN=64 -> launch_bounds(64) +BM=2, BN=64 -> launch_bounds(128) +BM=4, BN=64 -> launch_bounds(256) +BM=8, BN=64 -> launch_bounds(512) +``` + +host 侧仍需根据运行时 M/N/K 选择模板实例；这些 `if/else` 不会进入 +GPU kernel，也不会造成 wave divergence。 + +## 8. Fused Ticket 协议必须保持 + +每个 split CTA 写完 partial 后： + +```cpp +__syncthreads(); + +if (tid == 0) { + __threadfence(); + ticket = atomicAdd(&tile_done[tile_id], 1u); +} +``` + +顺序不能随意交换： + +1. 所有线程先写完当前 CTA 的 partial； +2. CTA barrier 确认 block 内写入已经发出； +3. device fence 发布 global partial； +4. 最后增加完成计数； +5. 获得 `split_k - 1` 票号的 CTA 执行 Reduce。 + +`atomicAdd` 只负责同一次 GEMM 的 splits。Per-stream workspace 的作用 +是防止不同 GEMM 调用共享同一个 counter 和 partial。 + +`is_last_split` 由 thread 0 写入 shared memory，并在 barrier 后供整个 +CTA 读取，因此它是 block-uniform 条件。非最后 CTA 可以在 ticket 后 +直接结束；只有最后 CTA 执行 Reduce： + +```cpp +if (is_last_split) { + if (row < M && col < N) { + // Reduce partial、scale、写 BF16 output。 + } + + __syncthreads(); + + if (tid == 0) + tile_done[tile_id] = 0u; +} +``` + +尾部不再需要第二次 `__threadfence()` 和 `atomicExch()`，理由是： + +1. 其他 stream 使用独立 counter； +2. 同一 stream 的下一次 kernel 必须等当前 kernel 完成； +3. 最后 ticket 出现时，其他 split CTA 已经完成对 counter 的最后一次访问； +4. last CTA 内的 barrier 保证所有输出线程先完成 store 指令，再由 thread 0 + 清零 counter。 + +不能据此删除 ticket 前的第一次 `__threadfence()`。第一次 fence 负责在 +发布完成计数前发布 global partial，是跨 CTA Reduce 正确性的必要条件。 + +## 9. 当前 Dispatch + +当前 split 数不变，但加入了 notebook 12 中已有交替 benchmark 证据的 +BM=2/BM=4 精确实例。 + +当前特殊实例： + +```cpp +if (M == 1 && N == 8192 && K == 1024) { + return launch_splitk_fused_instance<1, 64, 32>(...); +} + +const bool certified_small_bm_shape = + (K == 1024 && (N == 8192 || N == 4096)) || + (K == 4096 && N == 1024); + +if (M == 2 && certified_small_bm_shape) { + return launch_splitk_fused_instance<2, 64, 32>(...); +} + +if (M == 4 && certified_small_bm_shape) { + return launch_splitk_fused_instance<4, 64, 32>(...); +} + +return launch_splitk_fused_instance<8, 64, 32>(...); +``` + +不要仅根据 `M==2/4` 泛化到所有 shape。BM 同时影响计算 wave 和 +global-to-LDS 搬运并行度；当前只固化： + +```text +M=2/4, K=1024, N=8192 +M=2/4, K=1024, N=4096 +M=2/4, K=4096, N=1024 +``` + +后续仍可实验： + +```text +BK=64 +BK=128 +``` + +新增实例时应继续通过 `launch_splitk_fused_instance` +启动，不要重新复制 workspace 管理代码。 + +## 10. 并发语义 + +### 安全 + +- 同一显式 stream 连续调用； +- 多个显式 stream 并发调用； +- 多个 CPU 线程向同一显式 stream 提交； +- 同一进程内不同 HIP device 的显式或默认 stream； +- 不同 shape 在容量足够的 workspace 上复用。 + +### 需要注意 + +1. **默认 stream** + + 当前 key 已包含 device ID，可以区分不同 GPU 上的空 stream 句柄。 + 但如果传入 `nullptr` 且运行环境使用 per-thread default stream + 语义，同一 device 的不同 host 线程仍可能具有相同空句柄、却对应 + 不同执行序列。此时需要把 host thread 信息加入默认-stream key， + 或要求调用方传入显式 stream。 + +2. **Stream 生命周期** + + 当前 map 不知道外部 stream 何时销毁，因此不会主动释放对应设备 + workspace。适合 SGLang/PyTorch 中固定、长期存在的 stream。 + 若频繁创建和销毁临时 stream，应增加显式释放接口，或改成由 HIP + event 管理的有界 workspace slot pool。 + +3. **首次分配和扩容** + + 首次调用包含 `hipMalloc` 和 counter memset；扩容会同步对应 stream。 + benchmark 应区分 cold-start 与 steady-state。 + +## 11. Agent 后续实现建议 + +如果需要生产级完善，优先级如下： + +1. 明确处理 `nullptr` / per-thread default stream； +2. 为临时 stream 增加 workspace 回收机制； +3. 已知 shape 下可为每个 stream 预分配最大 workspace，避免热路径扩容； +4. 若 runtime 支持且验证稳定，可评估 `hipMallocAsync/hipFreeAsync`； +5. 保留 `(device_id, stream_handle)` 复合 key，不要退回单 stream key； +6. 不要用全局 `hipDeviceSynchronize` 代替 workspace 隔离； +7. 不要仅用 host mutex 保护单一 workspace——host mutex 在 launch 返回后 + 会释放，而不同 stream 的异步 kernel 仍可能重叠。 + +## 12. 验证清单 + +### 编译 + +- 使用目标 HIP 编译器； +- 目标架构为 gfx928； +- 检查 BM1/BN64 实例的 workgroup 上限为 64； +- 检查 BM8/BN64 实例的 workgroup 上限为 512。 + +### 单 stream + +- 连续运行不同 Split-K shape； +- 先小 shape、后大 shape，覆盖扩容路径； +- 先大 shape、后小 shape，覆盖容量复用； +- 与 CPU 或可信 GEMM reference 比较。 + +### 双 stream 正确性压力测试 + +建议两个 stream 使用不同输入模式，避免错误相互抵消： + +```text +stream A: + A/weight 填充模式 A + output A + +stream B: + A/weight 填充模式 B + output B +``` + +循环交错提交： + +```text +launch(A, streamA) +launch(B, streamB) +launch(A, streamA) +launch(B, streamB) +... +``` + +最后分别同步两个 stream，并验证所有输出。 + +至少覆盖： + +- 相同 shape、不同数据； +- 不同 shape； +- BM1 与 BM8 同时运行； +- 两个 stream 同时首次分配； +- 一个 stream 扩容时另一个 stream 正在计算。 + +### 多设备 + +- 在两个 HIP device 上分别使用显式 stream； +- 在两个 device 上分别使用默认 stream； +- 确认相同的空 stream 句柄映射到不同 `DeviceStreamKey`； +- 交错提交 Split-K，并分别与 reference 比较； +- 检查 workspace 的设备归属和 kernel 当前 device 一致。 + +### 性能 + +分别记录： + +- 首次调用； +- workspace 容量稳定后的调用； +- 单 stream； +- 双 stream 并发吞吐； +- `hipStreamSynchronize` 是否只在扩容时出现。 + +## 13. 当前验证状态 + +已完成： + +- 旧单例 `g_workspace/g_tile_done` 引用清理； +- 每-stream 指针接入 fused kernel launch； +- workspace key 扩展为 `(device_id, stream_handle)`； +- `hipGetDevice` 失败路径检查； +- host map mutex 与每-stream launch mutex 分层； +- 扩容前同-stream 同步； +- notebook 12 已认证 shape 的 BM=2/BM=4 精确 dispatch； +- 非最后 CTA 跳过 fused 尾部 barrier； +- 最后 CTA 使用普通 store 复位专属 counter，删除第二次 fence/atomic； +- 源码格式和引用静态检查。 + +尚未完成： + +- 当前环境没有 `hipcc`，未执行 gfx928 编译； +- 未执行实机双 stream 正确性压力测试； +- 未对 BM2/BM4 和 fused 尾部精简执行新旧交替 benchmark； +- 未验证默认 stream/per-thread default stream； +- 未执行实机多 device 并发压力测试； +- 未实现 stream 销毁后的 workspace 回收。 diff --git a/metainfer/tasks/opt_GEMM_kernel/notebooks/14_bandwith_opt.md b/metainfer/tasks/opt_GEMM_kernel/notebooks/14_bandwith_opt.md new file mode 100644 index 00000000..c37d0538 --- /dev/null +++ b/metainfer/tasks/opt_GEMM_kernel/notebooks/14_bandwith_opt.md @@ -0,0 +1,565 @@ +# 从旧 SDOT4 算子到 W8A8 packed Marlin/MMAC 的优化总结 + +本文专门解释目录中的算子相对以下旧实现做了什么改进： + +```text +/data/FF/MetaInfer/nodes/worker26/workspaces/ +opt-gemm-kernel-2067149b/010/submission/myGEMM_kernel.hip +``` + +目标读者是需要参考这些经验继续编写 gfx928 HIP kernel 的 AI 或工程师。 +旧版本对照是 `w8a8-marlin-fused_backup_20260730_pre_stripe`；最终生产证据来自 +`/home/FF/workspace/003/build.sh`、`w8a8_gemm_fused.cpp` 及其实际编译的 kernel +source。源码存在或进入最终共享库不等于生产启用；判断某个 shape 的真实路径时， +必须同时找到 dispatcher 的 shape guard 和对应 launch。 + +## 0. 最终 `/home/FF/workspace/003` 的证据边界 + +### 0.1 三种状态必须分开 + +- **生产 dispatch**：最终 `w8a8_gemm_fused.cpp` 对某个 shape family 有可达 guard， + 并调用对应 launcher。这是“最终版本实际启用”的必要证据。 +- **编译进最终库**：`build.sh` 将 source 编译进 `libw8a8_marlin_fused.so`，但某个 + kernel 仍可能只被特定 guard 使用，不能推广到其他 shape。 +- **实验记录**：workspace 中未进入最终 `build.sh`，或最终 dispatcher 没有可达调用的 + source/report，只能作为候选假设，不能写成生产 dispatch。 + +最终 `build.sh` 明确编译 packed-weight、small-M packed MMAC、BM64、N8192 BM32、 +两个 M32/K4096 winner、M4096 kernel 和最终 dispatcher。生产 route 由 dispatcher +进一步限定： + +| shape guard | 最终 launcher | 已启用的主要机制 | 主要代价/复核点 | +|---|---|---|---| +| `M<=32, N=4096, K=1024` | `launch_bm64_b_vgpr` | BM64；B-in-VGPR；A-only LDS；A direct-to-LDS；细粒度 `ds_read_b64`/MMAC interleave；bit-reverse N scheduling；stride-17 epilogue | B fragment 和 accumulator 增加 VGPR；重新检查 occupancy、HBM 和尾部 mask | +| `M<=32, N=8192, K=1024` | `launch_bm32_vgpr_db_n8192` | BM32；A direct-to-LDS；B global read 使用 SLC；A/B double buffering；B prefetch 与 MMAC 重叠 | SLC 和 prefetch 是该 shape 的实测选择，不应泛化；检查 cache、VGPR 和并发 wave | +| `M<=32, N=1536, K=4096` | `launch_m32_k4096_wqkv_winner` | shape-specific M32/K4096 pipeline | 只对该 N guard 有生产证据；与通用 split-K 比较总 dispatch operator time | +| `M<=32, N=512, K=4096` | `launch_m32_k4096_shared_winner` | shape-specific shared pipeline | LDS 容量、barrier 和 occupancy 必须重新实测 | +| `M=4096` 的 dispatcher 列举 N/K family | `launch_m4096_bm128_bn128_group4` | BM128xBN128xBK64；8 waves；`GROUP_M=4` CTA ordering 复用 L2 中的 B；`ds_read_b64` fragments feeding MMAC | tile 大、wave 多；检查边界 guard、寄存器/LDS 和不同 N/K 下的 L2 收益 | +| dispatcher 的其他明确 family | BM16/BN32、BM16/BN16 或 BM32/BN32 shared launcher | shape-specific tile 路由，避免一个通用 tile 覆盖所有 M/N/K | 每个 guard 独立验证，不能按相近 shape 推断 winner | +| 未命中上述专用路径的支持 small-M shape | packed MMAC split-K | packed W8A8；caller-owned persistent `partial`/`tile_done`；last-arriving CTA fused reduction | partial HBM、fence/atomic、split 数和 workspace 并发所有权 | + +因此，本 notebook 记录的是优化机制和已知适用 guard，不是要求 agent 复制最终 dispatch +表。新迭代必须从当前 `submission/` 和当前 frozen shapes 重新建立 route map，再用当前 +hipprof trace 的每次调用全部 dispatch `DurationNs` 之和比较。PMC 只解释 HBM、L2、 +register、LDS/scratch 和可靠可用的 wave/occupancy 信息；PMC replay duration 不是 latency。 +绝对微秒数和 workspace 中的历史 winner 名称都不是跨机器 SLA。 + +### 0.2 最终版本中确认启用的技术 + +| 技术 | 最终证据 | 适用范围 | 不应忽略的 trade-off | +|---|---|---|---| +| packed W8A8 MMAC | pack source、MMAC consumers、最终 dispatcher | 静态、可复用权重；多个 small/large-M route | packing 必须在计时区间外完成；layout/consumer lane mapping 必须一致 | +| B-in-VGPR、A-only LDS | `w8a8_bm64_b_vgpr_opt.cpp` 的生产 route | `M<=32,N=4096,K=1024` | VGPR 压力可能降低 resident waves | +| Direct-to-LDS | BM64、N8192 等生产 source | 对齐且 guard 满足的 A tile load | alignment、尾部和 builtin 语义必须保持安全 | +| A/B double buffering | N8192 production source；其他 route 以各自 source 为准 | memory latency 可与 MMAC 重叠的固定 shape | LDS/VGPR 增长和 barrier 次序可能抵消收益 | +| 细粒度 `ds_read`/MMAC interleave | BM64/M4096 production source | 有足够独立 fragment 的 MMAC loop | 调度改变必须以 operator time 和资源计数复核 | +| SLC read | N8192/K1024 production route | 当前只有该 guard 的启用证据 | cache policy 的收益依赖工作集和并发，不可全局开启 | +| bit-reverse N scheduling | BM64 production route | 当前 BM64 N-block ordering | 可能改变 L2 locality/负载均衡，需按 shape 复测 | +| stride-17/padded epilogue | BM64 production route | LDS epilogue transpose/store | 多占 LDS；必须验证 bank conflict 与 occupancy 的净效果 | +| shape-specific BM16/BM32/BM64 | dispatcher guards | dispatcher 明确列举的 family | 不能只按 M 或 N 相似就复用 | +| BM128xBN128 + `GROUP_M=4` | M4096 production route | dispatcher 明确列举的 M4096 family | 大 tile 资源压力和边界浪费 | +| caller-owned persistent split-K workspace | V2 query/launch ABI 与 fallback route | 需要 split-K 的支持 shape | 同一 workspace 不得被无保护的并发 stream 复用 | +| last-arriving-CTA fused reduction | packed small-M split-K consumer | partial/ticket protocol 使用的 route | device fence、ticket 清零、跨调用和并发正确性 | + +这些技术进入知识库的理由是最终 source 和 dispatch 证据，而不是历史报告里的单次 +speedup。任何后续采用都必须先提出一个有边界、可测、可回滚的假设，并以当前全部 +shape 的 hipprof operator time 及对应 PMC 证据决定保留或回滚。 + +## 1. 旧算子已经做了什么 + +旧 `myGEMM_kernel.hip` 并非朴素基线。它已经包含： + +- 128-bit global load； +- `__builtin_amdgcn_sdot4` INT8 点积； +- A/B LDS tiling； +- B 从 `[K,N]` 到 `[N,K]` 的运行时 LDS 转置； +- `BK_PAD=BK+4` 的 LDS padding； +- small-M Split-K； +- INT32 partial 后再做 FP32 scale 和 BF16 conversion； +- M>16 的 BM16/BN16/BK64 SDOT4 kernel。 + +旧 small-M 核心是“一线程负责一个输出元素”： + +```cpp +// myGEMM_kernel.hip +const int row = int(blockIdx.y) * BM + ty; +const int col = int(blockIdx.x) * BN + tx; +int32_t acc = 0; + +for (int kk = 0; kk < BK; kk += 4) { + const int32_t a_pack = + *reinterpret_cast(&a_tile[ty][kk]); + const int32_t b_pack = + *reinterpret_cast(&b_tile[tx][kk]); + acc = __builtin_amdgcn_sdot4(a_pack, b_pack, acc, false); +} +``` + +每处理一个 K tile，旧算子都要： + +```text +读取 A + 读取 raw W + 在 LDS 中转置 W + → __syncthreads() + → SDOT4 + → __syncthreads() + → 下一个 K tile +``` + +因此新版本的改进重点不是再加一层 vector load，而是改变计算 primitive、静态 +权重布局、同步结构、Split-K 归约协议和 workspace ABI。 + +## 2. 改进总览 + +| 方面 | 旧 myGEMM | packed Marlin/MMAC | 原理 | +|---|---|---|---| +| 计算 | 每线程一个 C，循环 SDOT4 | Wave64 执行 16×16×32 INT8 MMAC | 硬件矩阵指令复用 operand，减少点积/地址指令 | +| W 布局 | forward 读取 raw `[K,N]` | 模型加载期 pack 成 MMAC lane layout | 一次性预处理静态权重 | +| B 路径 | global→LDS transpose→VGPR | packed global→VGPR | 删除 B LDS、转置和 bank conflict | +| A staging | 每 K32 tile staging | 每 split staging 完整 A slice | 主 K 循环从每 tile 两个 barrier 降到一次 barrier | +| M tile | BM=1/2/4/8；M9..16 用 BM8 | 固定 BM16，非法行 mask | 同一 A slice 覆盖完整 M16 | +| Split-K reduce | 两次 kernel launch | last-arriving CTA 融合归约 | 减少短 kernel 的 launch 开销 | +| K 切分 | `span=K/split_k` | 按 K32 tile 比例切分 | 不要求等长整数 span，无遗漏和重复 | +| workspace | 库内 map/mutex/hipMalloc | 调用方查询并持久分配 | forward 不分配、不做 host mutex/sync | +| split_k | 常见 8/16/32 | shape 实测后常见 2/4/8 | 平衡并行度与 partial/reduce 成本 | +| M>16 | BM16/BN16/BK64 SDOT4 | BM16/BN64 packed MMAC no-split | 一次覆盖更多 N，消除运行时 B 转置 | + +## 3. 改进一：SDOT4 改成 Wave64 MMAC + +新算子使用 gfx928 的： + +```cpp +v_mmac_i32_16x16x32_i8 +``` + +关键封装： + +```cpp +typedef int int2_t __attribute__((ext_vector_type(2))); +typedef int int4_t __attribute__((ext_vector_type(4))); + +__device__ __forceinline__ int4_t mmac_i32_16x16x32_i8( + int2_t a, int2_t b, int4_t c) { + __builtin_amdgcn_sched_barrier(0); + __asm__ __volatile__( + "v_mmac_i32_16x16x32_i8 %0, %1, %2, %0" + : "+v"(c) : "v"(a), "v"(b)); + __builtin_amdgcn_sched_barrier(0); + return c; +} +``` + +正式 small-M tile 是： + +```text +BM=16, BN=64, BK=32 +4 Wave64/CTA，256 threads +wave 0..3 分别计算四个 N16 +row16 = lane & 15 +k_group = lane >> 4 +``` + +每 lane 提供 8B A 和 8B B，并持有四个 INT32 accumulator。MMAC accumulator +对应输出： + +```cpp +row = row_base + (lane & 15); +col = n_block * 64 + wave * 16 + (lane >> 4) + i * 4; // i=0..3 +``` + +原理：旧 SDOT4 仍由软件逐个输出组织点积；MMAC 把一个 Wave64 的数据组织交给 +矩阵指令，一次推进 16×16×32 tile，显著减少 dot/loop/address 指令。 + +AI 写新 kernel 时不能只复制 asm。必须同时保持 lane→A fragment、lane→B fragment、 +lane→C accumulator 三个映射完全一致。 + +## 4. 改进二：把 B 转置移到模型加载阶段 + +### 4.1 旧算子的重复成本 + +旧算子每个 forward、每个 K tile 都从 row-major W 读取 16 个连续 N，再写成 +K-contiguous LDS： + +```cpp +// old: global W[K,N] -> LDS b_tile[N][K_PAD] +Vec128 value = *reinterpret_cast(w + w_off); +#pragma unroll +for (int j = 0; j < 16; ++j) + b_tile[ln + j][lk] = reinterpret_cast(&value)[j]; +``` + +这会反复产生： + +- B LDS 空间； +- B transpose store 指令； +- LDS bank conflict； +- B 可消费前的 CTA barrier。 + +### 4.2 新 packed layout + +静态 W 在模型初始化时重排为： + +```text +packed[n_block64][k_tile32][n_group16][lane64][byte8] +``` + +精确映射： + +```cpp +kg = lane / 16; +col = lane % 16; +Ksrc = kt * 32 + kg * 8 + i; +Nsrc = nb * 64 + n_group * 16 + col; +packed[nb][kt][n_group][lane][i] = W[Ksrc][Nsrc]; +``` + +pack kernel 的关键代码： + +```cpp +const int n_group = tid >> 6; +const int lane = tid & 63; +const int k_group = lane >> 4; +const int col16 = lane & 15; +const int local_n = n_group * 16 + col16; + +const int2_t fragment = + load_int8x8(&weight_nk[local_n][k_group * 8]); +const size_t chunk = + (((size_t(n_block) * k_tile_count + k_tile) * 4 + n_group) + * 64 + lane); +reinterpret_cast(packed_weight)[chunk] = fragment; +``` + +forward 中每个 wave 的 64 lanes 读取连续 `64×8=512B`： + +```cpp +const size_t chunk = + (((size_t(n_block) * k_tile_count + global_tile) * 4 + wave) + * 64 + lane); +const int2_t b_fragment = + reinterpret_cast(packed_w)[chunk]; +``` + +这使 B 完全绕过 LDS，直接 global→VGPR→MMAC。packing 只适合可长期复用的静态 +推理权重；如果 W 每次调用都变化，必须把 packing 成本纳入端到端评价。 + +packed buffer 字节数： + +```text +ceil(N/64) * ceil(K/32) * 4 * 64 * 8 +``` + +raw 和 packed 指针不能互换，否则数值会错误但通常不会触发内存异常。 + +## 5. 改进三：whole-A-slice，删除 K 循环内 barrier + +旧 small-M kernel 每 K32 tile staging 一次 A/B，前后各一次 barrier。新 kernel +因为 B 不再使用 LDS，可以把一个 Split-K CTA 负责的全部 A 一次放进 LDS： + +```cpp +constexpr int kOptMaxTilesPerSplit = 16; +__shared__ __align__(16) +int8_t a_slice[kOptMaxTilesPerSplit][16][36]; + +// CTA cooperative 128-bit load all A tiles in this split +// ... +__syncthreads(); // 整个 slice 的主计算只需要这一次同步 +``` + +随后直接循环 LDS A 与 global packed B： + +```cpp +for (int local_tile = 0; local_tile < slice_tiles; ++local_tile) { + const int2_t a = load_int8x8( + &a_slice[local_tile][row16][k_group * 8]); + const int2_t b = load_b(local_tile); + acc = mmac_i32_16x16x32_i8(a, b, acc); +} +``` + +例如 K=4096、split_k=8 时，每 split 为 16 个 K32 tile。旧主循环约有 32 个 +barrier 点；新主计算只需一次 staging barrier。 + +`36=32+4` 的 stride 保留旧算子已经验证过的 padding 思路：改变相邻行的 LDS +bank 起点，同时保持 dword 对齐。 + +每 split 最多 16 个 K32 tile 不是算法限制，而是 LDS/occupancy 限制。dispatcher +必须选择足够大的 split_k，使每段不超过此上限。 + +## 6. 改进四：B prefetch 与双 MMAC accumulator + +新 kernel 在消费当前 B 前，提前发出下一 fragment 的 global load，并让偶/奇 K tile +写入不同 accumulator： + +```cpp +int4_t acc0{}, acc1{}; +int2_t b_next = load_b(0); + +for (int t = 0; t < slice_tiles; ++t) { + const int2_t b = b_next; + if (t + 1 < slice_tiles) + b_next = load_b(t + 1); + const int2_t a = load_int8x8(&a_slice[t][row16][k_group * 8]); + if ((t & 1) == 0) acc0 = mmac_i32_16x16x32_i8(a, b, acc0); + else acc1 = mmac_i32_16x16x32_i8(a, b, acc1); +} +acc = acc0 + acc1; +``` + +原理： + +- prefetch 让下一次 B memory latency 与当前 MMAC 尝试重叠； +- 两条 accumulator dependency chain 给调度器更多独立工作。 + +代价是增加 live VGPR。不要继续盲目增加 prefetch 深度或 accumulator 数；必须用 +hipprof 检查 VGPR、occupancy 和实际 latency。 + +## 7. 改进五：两次 launch 的 Split-K 改成单 kernel 融合归约 + +旧实现启动两个 kernel： + +```cpp +hipLaunchKernelGGL(small_m_splitk_sdot4_kernel, ...); +hipLaunchKernelGGL(reduce_splitk_scale_kernel, ...); +``` + +新实现保留 INT32 partial，但加入每个 `(m_block,n_block)` 的 ticket。每个 split 写完 +partial 后发布；最后到达的 CTA 负责归约、scale 和 BF16： + +```cpp +if (valid_output_lane) + __threadfence(); +__syncthreads(); + +if (tid == 0) { + const uint32_t old = atomicAdd(&tile_done[ticket_index], 1u); + is_last_split = (old == uint32_t(split_k - 1)); +} +__syncthreads(); + +if (is_last_split) { + // sum partial[s][row][col] + // float(sum) * a_scale[row] * w_scale[col] -> BF16 + __syncthreads(); + if (tid == 0) + tile_done[ticket_index] = 0u; +} +``` + +关键原理： + +- `__threadfence()` 必须发生在 atomic ticket 前，保证其他 CTA 能观察 partial； +- atomic 返回 `split_k-1` 的 CTA 是最后到达者； +- ticket 完成后清零，从而让持久 workspace 可用于下一次调用； +- 归约仍为 INT32，所以数值语义没有改变。 + +这样删除了独立 reduction GPU dispatch 及其 GPU 工作。当前任务的 hipprof operator +latency 不包含 host launch API 时间，但会包含 standalone reduction kernel 的 +`DurationNs`；因此是否获益必须看每次调用全部 GPU dispatch 的时长总和，而不是只看 +main MMAC。注意同一个 `tile_done` workspace 不可被多个并发 stream 无保护复用。 + +## 8. 改进六：外部持久 workspace 代替库内 map/mutex/hipMalloc + +旧算子在共享库内部维护 per-stream workspace： + +```cpp +static std::mutex g_ws_mutex; +static std::unordered_map, DeviceStreamKeyHash> g_workspaces; + +hipStreamSynchronize(stream); +hipFree(ws->partial); +hipMalloc(&ws->partial, required_bytes); +``` + +扩容时会同步 stream，且 launch path 涉及 host mutex 和隐藏状态。新 V2 ABI 改成： + +```cpp +query_w8a8_gemm_v2_workspace( + M, N, K, + &partial_elements, + &tile_done_elements, + &packed_weight_bytes, + &split_k); +``` + +调用者在初始化阶段分配并持久保存： + +```text +partial: split_k * M * N 个 int32 +tile_done: ceil(M/16) * ceil(N/64) 个 uint32 +packed W: 模型加载时生成一次 +``` + +forward 只传已有指针。收益是： + +- 无 forward-time `hipMalloc/hipFree`； +- 无库内全局 map 和 host mutex； +- workspace 生命周期和并发归调用方明确管理； +- server/runtime 可以提前规划显存。 + +## 9. 改进七：重新标定 Split-K + +旧 SDOT4 为增加 CTA 数，使用过较大的 split： + +```cpp +// old examples +if (K == 4096 && N == 1536) return M <= 8 ? 16 : 8; +if (K == 4096 && N == 1024) return M <= 4 ? 32 : 16; +if (K == 2048 && N == 4096) return M <= 8 ? 16 : 8; +``` + +MMAC CTA 每次完成更多计算，过大 split 会增加： + +- partial 写入量； +- fence/atomic 数量； +- reducer 读取量； +- 每 CTA 太短造成的效率损失。 + +因此新版本对关键 shape 扫描后常用 2/4/8： + +| N,K | M<=8 | M=16 | +|---|---:|---:| +| 1536,4096 | 8 | 8 | +| 8192,1024 | M1=4，其他=2 | 2 | +| 4096,2048 | 4 | 4 | +| 1024,4096 | 8 | 8 | +| 4096,512 | 2 | 2 | + +切分按 K32 tile 做： + +```cpp +begin = split * k_tile_count / split_k; +end = (split + 1) * k_tile_count / split_k; +``` + +相比旧 `span=K/split_k`，这种写法允许 tile 数不能被 split 整除，同时保证所有 K32 +tile 恰好被处理一次。 + +选择 split_k 的真实目标是平衡： + +```text +CTA 并行度收益 +vs. +短 CTA 效率 + partial HBM 流量 + fence/atomic + reduction +``` + +旧 SDOT4 的 split 表不能直接复制给 MMAC kernel。 + +## 10. M=9..16 和 M>16 的变化 + +旧算子： + +- M=9..16、Split-K 时固定使用 BM8，需要两个 M blocks； +- M=9..16、非 Split-K 时回退到 scalar GEMV； +- M>16 使用 BM16/BN16/BK64 SDOT4，并继续运行时 B LDS 转置。 + +新算子： + +- small-M 固定 BM16，M=9..16 一次覆盖，多余行自然 mask； +- pre-stripe 的 M>16 使用 BM16/BN64 packed MMAC no-split； +- no-split large path 无 partial、ticket 和 reducer，直接做 scale/BF16。 + +BN16→BN64 的意义是同一次 A staging 服务四个 N16 wave。配合 packed B，运行时不再 +转置 B。 + +但 pre-stripe large-M 仍会让不同 M16 CTA 重读同一 B。它不是最终 large-M 方案； +目录里的 stripe/shared-LDS/direct-LDS 实验就是继续探索跨 M block 的 B reuse。实验 +kernel 未必进入正式 `launch_w8a8_gemm_v2`，不可仅凭存在源码就当生产路径。 + +## 11. 哪些旧思路被保留 + +新版本并非推倒重来，保留了旧算子中正确的基础： + +- A/raw W 优先使用 128-bit aligned load，并保留 tail guard； +- K 基本 tile 仍为 32； +- LDS stride 使用 `+4` padding； +- Split-K partial 保持 INT32； +- scale 只在完整 INT32 accumulation 后执行； +- epilogue 仍为 FP32 scale 后 BF16 conversion； +- legacy raw-weight kernels 作为不能使用 packed 路径时的 fallback。 + +完整演进关系： + +```text +旧 SDOT4 tiled GEMM + + 保留 vector load / K32 / INT32 partial / BF16 epilogue + + 静态 W 离线 packing + + SDOT4 -> Wave64 MMAC + + B LDS transpose -> packed B direct VGPR load + + per-K32 A staging -> whole-split A slice + + separate reduce launch -> last-arriving-CTA fused reduce + + hidden workspace -> caller-owned persistent workspace += packed Marlin-fused kernel +``` + +## 12. 对 AI 最重要的实现约束 + +1. 不要把 raw `[K,N]` 指针传给 packed consumer。 +2. 修改 pack layout 时必须同步修改 consumer offset 和 lane mapping。 +3. `v_mmac` 的 A/B/C lane mapping必须先数学推导再编码。 +4. whole-A-slice 每 split 不得超过 16 个 K32 tile。 +5. 所有 CTA threads 必须以一致控制流到达 `__syncthreads()`。 +6. Split-K partial 发布必须先于 ticket atomic，并有 device-scope fence。 +7. `tile_done` 首次使用前清零；调用结束必须安全复位。 +8. 同一 workspace 不得被并发 stream 无保护复用。 +9. 增加 accumulator/prefetch 会提高 VGPR，必须检查 occupancy。 +10. 新 shape 必须扫描 split_k，而不是沿用旧 SDOT4 参数。 + +## 13. 验证与性能分析 + +构建： + +```bash +bash build.sh +``` + +正确性优先跑： + +```bash +python3 smoke_test_v2.py +python3 smoke_test_tp8_v2.py +``` + +reference 应采用 CPU INT32 GEMM，再做完全相同的 FP32 scale/BF16 conversion，要求 +逐元素 `mismatch == 0`。应覆盖 TP4/TP8、M=1/2/3/4/5/8/16/32 和尾部 shape。 + +性能最终使用当前任务冻结的 hipprof protocol： + +- 10 次 warmup dispatch group； +- 统计后续 100 次 operator call，trace 共 110 次； +- 每次调用把该 GEMM 的全部 GPU dispatch `DurationNs` 求和，再对最终 100 次取算术平均； +- 用 realtime host interval 和稳定 dispatch pattern 排除 preparation、packing、JIT 与无关 kernel； +- PMC/read/write replay 单独采集 HBM、L2、VGPR/AGPR/SGPR、LDS/scratch 等诊断， + replay duration 不进入 latency。 + +阶段归因可使用 `profile_phase_breakdown.py`： + +```text +main-only = MMAC + partial store +publish mode = main + fence/ticket +reduce-only = partial reduce + scale + BF16 +fused = 完整路径 +``` + +不要只优化 main MMAC。如果 publish/reducer 已占主要比例，应降低 split_k、改变 partial +layout 或 reducer mapping，而不是继续展开 MMAC。 + +## 14. 可直接交给 AI 的精简提示 + +```text +旧基线 myGEMM_kernel.hip 已有 128-bit load、SDOT4、A/B LDS tiling 和两阶段 +Split-K。新方案的结构性优化是：静态 W 在模型加载期 pack 成 +[N64][K32][N16-group][lane64][8B]；forward 中 B 直接 global->VGPR；使用 Wave64 +v_mmac_i32_16x16x32_i8，BM16/BN64/BK32、4 waves/CTA；每个 split 最多 16 个 +K32 tile，一次性 staging 完整 A slice 到 LDS [tiles][16][36]，主计算只做一次 CTA +barrier；预取下一 B fragment，偶/奇 tile 使用两条 accumulator chain。 + +Split-K 写 INT32 partial。有效 lanes threadfence 后，tid0 atomicAdd per-output-tile +ticket，最后 CTA 在同 kernel 中归约、scale、写 BF16并清 ticket，从而替代旧版第二个 +reduction launch。workspace 和 packed W 均由调用方初始化时持久分配，forward 不做 +hipMalloc/host mutex。split_k 必须按 MMAC 实测，不能照搬旧 SDOT4 的 16/32。 + +任何修改都必须保持 pack/consumer lane mapping、INT32 精确累加和 BF16 epilogue； +用 CPU INT32 reference 覆盖 TP4/TP8 与边界 M，并用 hipprof 检查 timing、HBM、L2、 +VGPR 和 LDS。 +``` diff --git a/metainfer/tasks/opt_GEMM_kernel/notebooks/README.md b/metainfer/tasks/opt_GEMM_kernel/notebooks/README.md index 89efda06..6197035c 100644 --- a/metainfer/tasks/opt_GEMM_kernel/notebooks/README.md +++ b/metainfer/tasks/opt_GEMM_kernel/notebooks/README.md @@ -32,7 +32,30 @@ remaining multi-stream workspace limitations, read `11_champion_engineering_DPP_alignment_generality.md`. For the measured follow-up work on large-M MMAC L2-aware CTA swizzling, -last-arriving-CTA fused split-K, the exact `M=1,N=8192,K=1024` BM=1 +last-arriving-CTA fused Split-K, the exact `M=1,N=8192,K=1024` BM=1 specialization, physical-versus-logical bandwidth interpretation, and the shape guards and concurrency cautions required to use those techniques, read `12_MMAC_CTA_swizzle_fused_splitK_BM1.md`. + +For workspace ownership and fused Split-K concurrency history, read +`13_stream_splitK workload_opt.md`. Its per-stream workspace design was not +fully compiled or stress-tested on K100 and is not the final caller-owned +workspace ABI, so use it as a correctness-risk record rather than a dispatch +prescription. + +For the final `/home/FF/workspace/003` packed W8A8 MMAC implementation and its +shape-specific bandwidth/compute techniques, read `14_bandwith_opt.md`. It +distinguishes production dispatches from experiments and records the evidence +source for each technique. + +## How agents must use this knowledge base + +These notes are historical evidence, constraints, and candidate hypotheses. +They are intentionally not exhaustive and never override the current +`submission/`, frozen shape contract, or fresh hipprof evidence. Before choosing +an optimization, inspect the current source and join each target shape with its +summed GPU operator time, dispatch breakdown, HBM/L2 metrics, and available +resource data. Derive a bounded, measurable, reversible hypothesis independently; +new ideas outside these notebooks are expected when the evidence supports them. +Absolute microsecond values in a notebook are machine/run-specific observations, +not portable performance requirements. diff --git a/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/__init__.py b/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/__init__.py index 523fdfa5..b9c263b3 100644 --- a/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/__init__.py +++ b/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/__init__.py @@ -2,7 +2,12 @@ from .champion import ChampionStore from .runner import EvaluationError, EvaluationResult, EvaluatorRunner -from .scoring import ScoreResult, compare_measurements, score_benchmark +from .scoring import ( + PromotionResult, + ScoreResult, + compare_against_champion, + compare_measurements, +) from .spec import BenchmarkCaseSpec, FrozenEvaluatorBundle, KernelTaskSpec, SpecError from .weights import FrozenWeightBundle @@ -15,8 +20,9 @@ "FrozenEvaluatorBundle", "FrozenWeightBundle", "KernelTaskSpec", + "PromotionResult", "ScoreResult", "SpecError", + "compare_against_champion", "compare_measurements", - "score_benchmark", ] diff --git a/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/champion.py b/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/champion.py index 31167978..1852232d 100644 --- a/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/champion.py +++ b/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/champion.py @@ -1,58 +1,156 @@ -"""Persistent champion/challenger selection for GEMM candidates.""" +"""Persistent per-shape Champion selection for GEMM candidates.""" from __future__ import annotations -import json import hashlib +import json import os import shutil import time from pathlib import Path -from typing import Any, Dict, Optional +from typing import Any, Dict, Optional, Sequence + +from .scoring import compare_against_champion + + +ReportReference = Dict[str, str] + + +def write_json_atomic(path: Path, data: Dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(path.suffix + ".tmp") + tmp.write_text(json.dumps(data, indent=2), encoding="utf-8") + os.replace(tmp, path) + + +def make_report_reference(state_root: Path, report_path: Path) -> ReportReference: + root = state_root.resolve() + path = report_path.resolve() + try: + relative = path.relative_to(root) + except ValueError as exc: + raise RuntimeError(f"performance report is outside task state: {path}") from exc + if not path.is_file(): + raise RuntimeError(f"performance report is missing: {path}") + return { + "path": relative.as_posix(), + "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + } + + +def load_report_reference( + state_root: Path, reference: Dict[str, Any], +) -> Dict[str, Any]: + root = state_root.resolve() + relative = str(reference.get("path") or "").strip() + expected = str(reference.get("sha256") or "").strip() + if not relative or not expected: + raise RuntimeError("performance report reference is incomplete") + path = (root / relative).resolve() + try: + path.relative_to(root) + except ValueError as exc: + raise RuntimeError("performance report reference escapes task state") from exc + try: + payload = path.read_bytes() + except OSError as exc: + raise RuntimeError(f"performance report is missing: {relative}") from exc + actual = hashlib.sha256(payload).hexdigest() + if actual != expected: + raise RuntimeError( + f"performance report changed: expected {expected}, got {actual}" + ) + try: + report = json.loads(payload) + except json.JSONDecodeError as exc: + raise RuntimeError(f"performance report is invalid: {relative}") from exc + if not isinstance(report, dict): + raise RuntimeError(f"performance report must be an object: {relative}") + return report + + +def champion_report_reference( + state_root: Path, record: Dict[str, Any], +) -> ReportReference: + reference = record.get("measurement_report") + if isinstance(reference, dict): + load_report_reference(state_root, reference) + return {"path": str(reference["path"]), "sha256": str(reference["sha256"])} + + kind = str(record.get("kind") or "hip") + iteration = int(record.get("iteration") or 0) + if kind == "triton": + path = state_root / "baseline" / "baseline-benchmark-report.json" + elif iteration == 0: + path = ( + state_root / "certified" / "initial-hip" + / "candidate-benchmark-report.json" + ) + else: + path = state_root / "logs" / f"{iteration:03d}" / "candidate-benchmark-report.json" + return make_report_reference(state_root, path) class ChampionStore: - def __init__(self, root: Path, noise_threshold: float) -> None: + def __init__( + self, + root: Path, + noise_threshold: float, + expected_case_ids: Sequence[str], + ) -> None: self.root = root + self.state_root = root.parent self.submission_dir = root / "submission" self.record_path = root / "champion.json" self.noise_threshold = noise_threshold + self.expected_case_ids = list(expected_case_ids) def initialize(self, initial_submission: Optional[Path]) -> None: + """Legacy entry point retained for callers that seed a HIP Champion.""" if self.record_path.exists(): self.load() return + if initial_submission is None: + raise RuntimeError("initial HIP Champion requires a submission") + report_path = ( + self.state_root / "certified" / "initial-hip" + / "candidate-benchmark-report.json" + ) + reference = make_report_reference(self.state_root, report_path) self.root.mkdir(parents=True, exist_ok=True) - if initial_submission and initial_submission.is_dir(): - shutil.copytree(initial_submission, self.submission_dir, dirs_exist_ok=True) + shutil.copytree(initial_submission, self.submission_dir, dirs_exist_ok=True) self._write({ + "schema_version": 2, "kind": "hip", "iteration": 0, - "weighted_speedup": 1.0, + "measurement_report": reference, "submission_sha256": _tree_digest(self.submission_dir), "promoted_at": time.time(), - "reason": "initial baseline", + "reason": "initial HIP Champion", }) - def initialize_triton(self, certified_baseline: Dict[str, Any]) -> None: - """Initialize the arena incumbent from the frozen Triton evaluation.""" + def initialize_triton(self, measurement_report: ReportReference) -> None: + """Initialize the arena incumbent from frozen Triton measurements.""" if self.record_path.exists(): self.load() return + load_report_reference(self.state_root, measurement_report) self.root.mkdir(parents=True, exist_ok=True) self._write({ + "schema_version": 2, "kind": "triton", "iteration": 0, - "weighted_speedup": 1.0, - "baseline_manifest_sha256": certified_baseline.get("manifest_sha256"), + "measurement_report": dict(measurement_report), "promoted_at": time.time(), "reason": "certified Triton baseline", }) def load(self) -> Dict[str, Any]: if not self.record_path.exists(): - return {"iteration": 0, "weighted_speedup": 1.0} + return {"schema_version": 2, "kind": "triton", "iteration": 0} record = json.loads(self.record_path.read_text(encoding="utf-8")) + if not isinstance(record, dict): + raise RuntimeError("champion record must be an object") if record.get("kind", "hip") == "hip": expected = record.get("submission_sha256") if not expected or not self.submission_dir.is_dir(): @@ -60,25 +158,31 @@ def load(self) -> Dict[str, Any]: actual = _tree_digest(self.submission_dir) if actual != expected: raise RuntimeError("champion submission changed outside promotion") - return record + reference = champion_report_reference(self.state_root, record) + load_report_reference(self.state_root, reference) + return {**record, "measurement_report": reference} def consider( self, iteration: int, candidate_dir: Path, - score: Dict[str, Any], + candidate_report: ReportReference, + same_round_incumbent_report: ReportReference, ) -> tuple[bool, str, Dict[str, Any]]: current = self.load() - candidate_speedup = float(score.get("weighted_speedup", 0.0)) - current_speedup = float(current.get("weighted_speedup", 1.0)) - if not bool(score.get("passed")): - return False, "acceptance gates failed", current - required = current_speedup * (1.0 + self.noise_threshold) - if candidate_speedup < required: - return False, ( - f"speedup {candidate_speedup:.6f} did not beat champion " - f"{current_speedup:.6f} by noise threshold {self.noise_threshold:.2%}" - ), current + candidate = load_report_reference(self.state_root, candidate_report) + incumbent = load_report_reference( + self.state_root, same_round_incumbent_report + ) + promotion_gate = compare_against_champion( + incumbent.get("cases") or [], + candidate.get("cases") or [], + self.expected_case_ids, + self.noise_threshold, + strict=True, + ) + if not promotion_gate.passed: + return False, "; ".join(promotion_gate.reasons), current replacement = self.root / "submission.next" if replacement.exists(): @@ -88,22 +192,20 @@ def consider( shutil.rmtree(self.submission_dir) os.replace(replacement, self.submission_dir) record = { + "schema_version": 2, "kind": "hip", "iteration": iteration, - "weighted_speedup": candidate_speedup, - "critical_regression": float(score.get("critical_regression", 0.0)), + "measurement_report": dict(candidate_report), + "promotion_incumbent_report": dict(same_round_incumbent_report), "submission_sha256": _tree_digest(self.submission_dir), "promoted_at": time.time(), - "reason": "candidate passed all gates and beat the current champion", + "reason": "every shape beat the same-round Champion hipprof trace beyond the noise gate", } self._write(record) return True, record["reason"], record def _write(self, data: Dict[str, Any]) -> None: - self.root.mkdir(parents=True, exist_ok=True) - tmp = self.record_path.with_suffix(".tmp") - tmp.write_text(json.dumps(data, indent=2), encoding="utf-8") - os.replace(tmp, self.record_path) + write_json_atomic(self.record_path, data) def _tree_digest(root: Path) -> str: diff --git a/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/runner.py b/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/runner.py index 9ae2dc35..fc26ceaf 100644 --- a/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/runner.py +++ b/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/runner.py @@ -132,6 +132,30 @@ def run( report["build_fingerprint"] = build_fingerprint return self._validate(phase, report, role=role, baseline_report=baseline_report) + def validate_benchmark_report( + self, + report: Dict[str, Any], + *, + role: str, + build_fingerprint: str, + baseline_report: Optional[Dict[str, Any]] = None, + ) -> EvaluationResult: + """Validate hipprof measurements without executing a timing command.""" + if role not in {"baseline", "candidate"}: + raise EvaluationError(f"invalid evaluation role: {role}") + self.bundle.verify() + if self.private_verifier is not None: + self.private_verifier() + normalized = dict(report) + normalized["evaluation_role"] = role + normalized["build_fingerprint"] = build_fingerprint + return self._validate( + "benchmark", + normalized, + role=role, + baseline_report=baseline_report, + ) + def _validate( self, phase: str, diff --git a/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/scoring.py b/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/scoring.py index 4b941814..8a85f6d6 100644 --- a/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/scoring.py +++ b/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/scoring.py @@ -1,4 +1,4 @@ -"""Deterministic multi-shape scoring and promotion gates.""" +"""Deterministic per-shape benchmark and promotion gates.""" from __future__ import annotations @@ -12,8 +12,8 @@ @dataclass(frozen=True) class ScoreResult: passed: bool - weighted_speedup: float - critical_regression: float + worst_case_speedup: float + failed_case_ids: List[str] = field(default_factory=list) missing_case_ids: List[str] = field(default_factory=list) reasons: List[str] = field(default_factory=list) cases: List[Dict[str, Any]] = field(default_factory=list) @@ -22,79 +22,17 @@ def to_dict(self) -> Dict[str, Any]: return asdict(self) -def score_benchmark( - cases: Sequence[Dict[str, Any]], - expected_case_ids: Sequence[str], - acceptance: AcceptanceSpec, -) -> ScoreResult: - by_id: Dict[str, Dict[str, Any]] = {} - reasons: List[str] = [] - normalized: List[Dict[str, Any]] = [] - for raw in cases: - case_id = str(raw.get("id") or "").strip() - if not case_id or case_id in by_id: - reasons.append(f"invalid or duplicate benchmark case id: {case_id!r}") - continue - try: - baseline_ms = float(raw["baseline_ms"]) - candidate_ms = float(raw["candidate_ms"]) - weight = float(raw.get("weight", 1.0)) - except (KeyError, TypeError, ValueError): - reasons.append(f"case {case_id!r} has invalid timing fields") - continue - if ( - not math.isfinite(baseline_ms) - or not math.isfinite(candidate_ms) - or not math.isfinite(weight) - or baseline_ms <= 0 - or candidate_ms <= 0 - or weight <= 0 - ): - reasons.append(f"case {case_id!r} timings and weight must be positive") - continue - item = { - "id": case_id, - "baseline_ms": baseline_ms, - "candidate_ms": candidate_ms, - "weight": weight, - "critical": bool(raw.get("critical", False)), - "speedup": baseline_ms / candidate_ms, - "regression": candidate_ms / baseline_ms - 1.0, - } - by_id[case_id] = item - normalized.append(item) - - missing = sorted(set(expected_case_ids) - set(by_id)) - unexpected = sorted(set(by_id) - set(expected_case_ids)) - if missing and acceptance.require_all_cases: - reasons.append(f"missing benchmark cases: {missing}") - if unexpected: - reasons.append(f"unexpected benchmark cases: {unexpected}") - - expected = [by_id[cid] for cid in expected_case_ids if cid in by_id] - base_work = sum(item["weight"] * item["baseline_ms"] for item in expected) - candidate_work = sum(item["weight"] * item["candidate_ms"] for item in expected) - weighted = base_work / candidate_work if candidate_work > 0 else 0.0 - critical = [item["regression"] for item in expected if item["critical"]] - worst_critical = max(critical, default=0.0) +@dataclass(frozen=True) +class PromotionResult: + passed: bool + noise_threshold: float + failed_case_ids: List[str] = field(default_factory=list) + missing_case_ids: List[str] = field(default_factory=list) + reasons: List[str] = field(default_factory=list) + cases: List[Dict[str, Any]] = field(default_factory=list) - if weighted < acceptance.min_weighted_speedup: - reasons.append( - f"weighted speedup {weighted:.6f} < minimum {acceptance.min_weighted_speedup:.6f}" - ) - if worst_critical > acceptance.max_critical_regression: - reasons.append( - f"critical regression {worst_critical:.2%} exceeds " - f"{acceptance.max_critical_regression:.2%}" - ) - return ScoreResult( - passed=not reasons, - weighted_speedup=weighted, - critical_regression=worst_critical, - missing_case_ids=missing, - reasons=reasons, - cases=normalized, - ) + def to_dict(self) -> Dict[str, Any]: + return asdict(self) def compare_measurements( @@ -103,11 +41,8 @@ def compare_measurements( case_specs: Sequence[BenchmarkCaseSpec], acceptance: AcceptanceSpec, ) -> ScoreResult: - """Compare independent baseline/candidate measurements. - - Weights and criticality come from the frozen task spec, never from either - measurement report. This prevents per-iteration workload drift. - """ + """Require every candidate shape to beat the frozen Triton measurement.""" + del acceptance # Every declared shape is an unconditional hard gate. baseline, baseline_errors = _measurement_map(baseline_cases, "baseline") candidate, candidate_errors = _measurement_map(candidate_cases, "candidate") expected = [case.id for case in case_specs] @@ -116,23 +51,25 @@ def compare_measurements( ) unexpected = sorted((set(baseline) | set(candidate)) - set(expected)) reasons = [*baseline_errors, *candidate_errors] - if missing and acceptance.require_all_cases: + if missing: reasons.append(f"missing benchmark cases: {missing}") if unexpected: reasons.append(f"unexpected benchmark cases: {unexpected}") + failed: List[str] = [] normalized: List[Dict[str, Any]] = [] for spec in case_specs: if spec.id not in baseline or spec.id not in candidate: continue base_ms = baseline[spec.id] cand_ms = candidate[spec.id] + speedup = base_ms / cand_ms + if cand_ms >= base_ms: + failed.append(spec.id) normalized.append({ "id": spec.id, "baseline_ms": base_ms, "candidate_ms": cand_ms, - "weight": spec.weight, - "critical": spec.critical, "shape": spec.shape, "flops": spec.flops, "bytes": spec.bytes, @@ -140,32 +77,77 @@ def compare_measurements( "candidate_tflops": _rate(spec.flops, cand_ms, 1e9), "baseline_bandwidth_gbps": _rate(spec.bytes, base_ms, 1e6), "candidate_bandwidth_gbps": _rate(spec.bytes, cand_ms, 1e6), - "speedup": base_ms / cand_ms, + "speedup": speedup, "regression": cand_ms / base_ms - 1.0, + "passed": cand_ms < base_ms, }) + if failed: + reasons.append(f"candidate did not beat baseline for cases: {failed}") + worst = min((case["speedup"] for case in normalized), default=0.0) + return ScoreResult( + passed=not reasons, + worst_case_speedup=worst, + failed_case_ids=failed, + missing_case_ids=missing, + reasons=reasons, + cases=normalized, + ) + - base_work = sum(item["weight"] * item["baseline_ms"] for item in normalized) - candidate_work = sum(item["weight"] * item["candidate_ms"] for item in normalized) - weighted = base_work / candidate_work if candidate_work > 0 else 0.0 - worst_critical = max( - (item["regression"] for item in normalized if item["critical"]), default=0.0 +def compare_against_champion( + champion_cases: Sequence[Dict[str, Any]], + candidate_cases: Sequence[Dict[str, Any]], + expected_case_ids: Sequence[str], + noise_threshold: float, + *, + strict: bool = False, +) -> PromotionResult: + """Require every shape to improve on Champion beyond the noise floor.""" + champion, champion_errors = _measurement_map(champion_cases, "champion") + candidate, candidate_errors = _measurement_map(candidate_cases, "candidate") + expected = list(expected_case_ids) + missing = sorted( + (set(expected) - set(champion)) | (set(expected) - set(candidate)) ) - if weighted < acceptance.min_weighted_speedup: - reasons.append( - f"weighted speedup {weighted:.6f} < minimum {acceptance.min_weighted_speedup:.6f}" - ) - if worst_critical > acceptance.max_critical_regression: + unexpected = sorted((set(champion) | set(candidate)) - set(expected)) + reasons = [*champion_errors, *candidate_errors] + if missing: + reasons.append(f"missing champion comparison cases: {missing}") + if unexpected: + reasons.append(f"unexpected champion comparison cases: {unexpected}") + + failed: List[str] = [] + comparisons: List[Dict[str, Any]] = [] + for case_id in expected: + if case_id not in champion or case_id not in candidate: + continue + champion_ms = champion[case_id] + candidate_ms = candidate[case_id] + required_ms = champion_ms * (1.0 - noise_threshold) + passed = candidate_ms < required_ms if strict else candidate_ms <= required_ms + if not passed: + failed.append(case_id) + comparisons.append({ + "id": case_id, + "champion_ms": champion_ms, + "candidate_ms": candidate_ms, + "required_ms": required_ms, + "speedup_vs_champion": champion_ms / candidate_ms, + "improvement": 1.0 - candidate_ms / champion_ms, + "passed": passed, + }) + if failed: reasons.append( - f"critical regression {worst_critical:.2%} exceeds " - f"{acceptance.max_critical_regression:.2%}" + "candidate did not beat champion beyond noise threshold for cases: " + f"{failed}" ) - return ScoreResult( + return PromotionResult( passed=not reasons, - weighted_speedup=weighted, - critical_regression=worst_critical, + noise_threshold=noise_threshold, + failed_case_ids=failed, missing_case_ids=missing, reasons=reasons, - cases=normalized, + cases=comparisons, ) @@ -181,6 +163,9 @@ def _measurement_map( values: Dict[str, float] = {} errors: List[str] = [] for raw in cases: + if not isinstance(raw, dict): + errors.append(f"{label} benchmark case must be an object") + continue case_id = str(raw.get("id") or "").strip() if not case_id or case_id in values: errors.append(f"invalid or duplicate {label} case id: {case_id!r}") diff --git a/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/spec.py b/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/spec.py index fa81c2ab..b80b6a50 100644 --- a/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/spec.py +++ b/metainfer/tasks/opt_GEMM_kernel/orchestrator/evaluator/spec.py @@ -18,8 +18,16 @@ class SpecError(ValueError): pass -_PHASES = ("correctness", "benchmark") -_EXTRA_COMMANDS = ("profile",) +_REQUIRED_COMMANDS = ("correctness", "profile") +_REQUIRED_BENCHMARK_METHOD = { + "timer": "hipprof_gpu_kernel_duration_ns", + "statistic": "arithmetic_mean", + "operator_aggregation": "sum_gpu_kernel_duration_per_call", + "synchronization": "hipprof_trace", + "timed_scope": "operator_gpu_dispatches_only", + "host_launch_time_included": False, + "pmc_timing_used": False, +} @dataclass(frozen=True) @@ -30,17 +38,12 @@ class CommandSpec: @dataclass(frozen=True) class AcceptanceSpec: - min_weighted_speedup: float = 1.0 noise_threshold: float = 0.01 - max_critical_regression: float = 0.03 - require_all_cases: bool = True @dataclass(frozen=True) class BenchmarkCaseSpec: id: str - weight: float = 1.0 - critical: bool = False shape: Optional[Dict[str, int]] = None flops: Optional[float] = None bytes: Optional[float] = None @@ -78,7 +81,7 @@ def load(cls, path: Path) -> "KernelTaskSpec": if not isinstance(commands_raw, dict): raise SpecError("task.yaml requires commands mapping") commands: Dict[str, CommandSpec] = {} - for phase in _PHASES: + for phase in _REQUIRED_COMMANDS: item = commands_raw.get(phase) if not isinstance(item, dict): raise SpecError(f"commands.{phase} must be a mapping") @@ -89,14 +92,6 @@ def load(cls, path: Path) -> "KernelTaskSpec": if timeout_s < 1 or timeout_s > 86_400: raise SpecError(f"commands.{phase}.timeout_s must be in [1, 86400]") commands[phase] = CommandSpec(list(argv), timeout_s) - for phase in _EXTRA_COMMANDS: - item = commands_raw.get(phase) - if isinstance(item, dict): - argv = item.get("argv") - if isinstance(argv, list) and argv and all(isinstance(v, str) and v for v in argv): - timeout_s = int(item.get("timeout_s", 600)) - if 1 <= timeout_s <= 86_400: - commands[phase] = CommandSpec(list(argv), timeout_s) cases = raw.get("cases") or {} if not isinstance(cases, dict): @@ -118,12 +113,32 @@ def load(cls, path: Path) -> "KernelTaskSpec": try: warmup = int(protocol["warmup"]) samples = int(protocol["samples"]) + trace_calls = int(protocol["trace_calls"]) timer = str(protocol["timer"]).strip() except (KeyError, TypeError, ValueError) as exc: - raise SpecError("benchmark_protocol requires warmup, samples, and timer") from exc - if warmup < 1 or samples < 3 or not timer: - raise SpecError("benchmark protocol requires warmup>=1, samples>=3, and timer") - protocol = {**protocol, "warmup": warmup, "samples": samples, "timer": timer} + raise SpecError( + "benchmark_protocol requires warmup, samples, trace_calls, and timer" + ) from exc + if warmup < 0 or samples < 3 or trace_calls != warmup + samples or not timer: + raise SpecError( + "benchmark protocol requires samples>=3 and trace_calls=warmup+samples" + ) + for key, expected_value in _REQUIRED_BENCHMARK_METHOD.items(): + actual_value = protocol.get(key) + if actual_value != expected_value or ( + isinstance(expected_value, bool) + and type(actual_value) is not bool + ): + raise SpecError( + f"benchmark_protocol.{key} must be {expected_value!r}" + ) + protocol = { + **protocol, + "warmup": warmup, + "samples": samples, + "trace_calls": trace_calls, + "timer": timer, + } private = _unique_ids(cases.get("private", []), "cases.private", allow_empty=True) unknown_private = sorted(set(private) - set(correctness)) if unknown_private: @@ -133,17 +148,10 @@ def load(cls, path: Path) -> "KernelTaskSpec": if not isinstance(acc_raw, dict): raise SpecError("acceptance must be a mapping") acceptance = AcceptanceSpec( - min_weighted_speedup=float(acc_raw.get("min_weighted_speedup", 1.0)), noise_threshold=float(acc_raw.get("noise_threshold", 0.01)), - max_critical_regression=float(acc_raw.get("max_critical_regression", 0.03)), - require_all_cases=bool(acc_raw.get("require_all_cases", True)), ) - if not math.isfinite(acceptance.min_weighted_speedup) or acceptance.min_weighted_speedup <= 0: - raise SpecError("min_weighted_speedup must be positive") if not math.isfinite(acceptance.noise_threshold) or not 0 <= acceptance.noise_threshold < 1: raise SpecError("noise_threshold must be in [0, 1)") - if not math.isfinite(acceptance.max_critical_regression) or not 0 <= acceptance.max_critical_regression < 1: - raise SpecError("max_critical_regression must be in [0, 1)") return cls( name=name, public_contract=public_contract, @@ -167,8 +175,6 @@ def agent_contract(self) -> Dict[str, Any]: { "id": case.id, "shape": case.shape, - "weight": case.weight, - "critical": case.critical, } for case in self.benchmark_cases if case.id not in private @@ -220,12 +226,8 @@ def _benchmark_cases(value: Any) -> List[BenchmarkCaseSpec]: case = BenchmarkCaseSpec(item) elif isinstance(item, dict): case_id = str(item.get("id") or "").strip() - try: - weight = float(item.get("weight", 1.0)) - except (TypeError, ValueError) as exc: - raise SpecError(f"benchmark case {case_id!r} has invalid weight") from exc - if not case_id or not math.isfinite(weight) or weight <= 0: - raise SpecError("benchmark case id and positive weight are required") + if not case_id: + raise SpecError("benchmark case id is required") shape = _benchmark_shape(item, case_id) if shape is None: raise SpecError( @@ -235,17 +237,15 @@ def _benchmark_cases(value: Any) -> List[BenchmarkCaseSpec]: transferred = _optional_positive_number( item.get("bytes"), f"benchmark case {case_id!r} bytes" ) - if flops is None and shape is not None: + if flops is None: flops = float( 2 * shape["m"] * shape["n"] * shape["k"] * shape["batch"] ) case = BenchmarkCaseSpec( - case_id, - weight, - bool(item.get("critical", False)), - shape, - flops, - transferred, + id=case_id, + shape=shape, + flops=flops, + bytes=transferred, ) else: raise SpecError("benchmark cases must be strings or mappings") @@ -266,21 +266,10 @@ def _benchmark_matrix(value: Dict[str, Any]) -> List[BenchmarkCaseSpec]: raise SpecError("cases.benchmark mapping requires matrix") try: m_values = [int(item) for item in matrix["m_values"]] - large_m = int(matrix.get("large_m", 4096)) - small_total = float(matrix.get("small_m_total_weight", 0.5)) - large_weight = float(matrix.get("large_m_weight", 0.5)) except (KeyError, TypeError, ValueError) as exc: - raise SpecError("benchmark matrix has invalid M values or weights") from exc + raise SpecError("benchmark matrix has invalid M values") from exc if not m_values or len(set(m_values)) != len(m_values) or any(m <= 0 for m in m_values): raise SpecError("benchmark matrix m_values must be unique positive integers") - small_values = [m for m in m_values if m != large_m] - if large_m not in m_values or not small_values: - raise SpecError("benchmark matrix must contain large_m and at least one small M") - if not math.isfinite(small_total) or small_total <= 0: - raise SpecError("benchmark matrix small_m_total_weight must be positive") - if not math.isfinite(large_weight) or large_weight <= 0: - raise SpecError("benchmark matrix large_m_weight must be positive") - critical_m = {int(item) for item in matrix.get("critical_m", [1, large_m])} workloads = matrix.get("workloads") if not isinstance(workloads, list) or not workloads: raise SpecError("benchmark matrix workloads must be a non-empty list") @@ -294,22 +283,16 @@ def _benchmark_matrix(value: Dict[str, Any]) -> List[BenchmarkCaseSpec]: n = int(workload["n"]) k = int(workload["k"]) batch = int(workload.get("batch", 1)) - workload_weight = float(workload.get("weight", 1.0)) except (KeyError, TypeError, ValueError) as exc: raise SpecError(f"benchmark matrix workload {workload_id!r} is invalid") from exc - if not workload_id or min(n, k, batch) <= 0 or workload_weight <= 0: + if not workload_id or min(n, k, batch) <= 0: raise SpecError(f"benchmark matrix workload {workload_id!r} is invalid") for m in m_values: shape = {"m": m, "n": n, "k": k, "batch": batch} - weight = workload_weight * ( - large_weight if m == large_m else small_total / len(small_values) - ) transferred = float(batch * (m * k + k * n + 4 * m + 4 * n + 2 * m * n)) cases.append( BenchmarkCaseSpec( id=f"{workload_id}-m{m}", - weight=weight, - critical=m in critical_m, shape=shape, flops=float(2 * m * n * k * batch), bytes=transferred, diff --git a/metainfer/tasks/opt_GEMM_kernel/orchestrator/hardware_profiles.yaml b/metainfer/tasks/opt_GEMM_kernel/orchestrator/hardware_profiles.yaml index ea9070e5..8e2b53e4 100644 --- a/metainfer/tasks/opt_GEMM_kernel/orchestrator/hardware_profiles.yaml +++ b/metainfer/tasks/opt_GEMM_kernel/orchestrator/hardware_profiles.yaml @@ -29,13 +29,7 @@ profiles: required: true tool_candidates: - /opt/dtk/bin/hipprof - - /opt/dtk/rocprofiler/bin/rocprof - - /opt/dtk/bin/rocprofv3 - - /opt/dtk/bin/rocprof - - /opt/rocm/bin/rocprofv3 - - /opt/rocm/bin/rocprof - - rocprofv3 - - rocprof + - hipprof representative_cases: - wq-b-tp4-m1 - wq-b-tp4-m16 diff --git a/metainfer/tasks/opt_GEMM_kernel/orchestrator/iteration_record.py b/metainfer/tasks/opt_GEMM_kernel/orchestrator/iteration_record.py index 73def414..ca06bfe1 100644 --- a/metainfer/tasks/opt_GEMM_kernel/orchestrator/iteration_record.py +++ b/metainfer/tasks/opt_GEMM_kernel/orchestrator/iteration_record.py @@ -18,7 +18,10 @@ class IterationRecord: failure_reason: Optional[str] = None phases: Dict[str, Dict[str, Any]] = field(default_factory=dict) score: Dict[str, Any] = field(default_factory=dict) - hardware_profile: Dict[str, Any] = field(default_factory=dict) + measurement_report: Dict[str, str] = field(default_factory=dict) + profile_report: Dict[str, str] = field(default_factory=dict) + incumbent_measurement_report: Dict[str, str] = field(default_factory=dict) + incumbent_profile_report: Dict[str, str] = field(default_factory=dict) promoted: bool = False champion_iteration: int = 0 artifacts: List[str] = field(default_factory=list) diff --git a/metainfer/tasks/opt_GEMM_kernel/orchestrator/orchestrator.py b/metainfer/tasks/opt_GEMM_kernel/orchestrator/orchestrator.py index d9b0e2f4..7839ddd0 100644 --- a/metainfer/tasks/opt_GEMM_kernel/orchestrator/orchestrator.py +++ b/metainfer/tasks/opt_GEMM_kernel/orchestrator/orchestrator.py @@ -89,6 +89,7 @@ def run_with_requirements( "METAINFER_WEIGHT_SHA256": weight_bundle.digest, }, harness_argv=harness_argv, + benchmark_protocol=bundle.spec.benchmark_protocol, ) initial_value = str(req_field(req, "initial_submission") or "").strip() initial_submission = Path(initial_value).expanduser().resolve() if initial_value else None diff --git a/metainfer/tasks/opt_GEMM_kernel/orchestrator/pipeline.py b/metainfer/tasks/opt_GEMM_kernel/orchestrator/pipeline.py index 70e31336..c69dc6d7 100644 --- a/metainfer/tasks/opt_GEMM_kernel/orchestrator/pipeline.py +++ b/metainfer/tasks/opt_GEMM_kernel/orchestrator/pipeline.py @@ -5,6 +5,7 @@ import json import hashlib import shutil +import statistics import time from dataclasses import dataclass, field from pathlib import Path @@ -23,6 +24,12 @@ FrozenEvaluatorBundle, FrozenWeightBundle, ) +from .evaluator.champion import ( + ReportReference, + load_report_reference, + make_report_reference, + write_json_atomic, +) from .guidance import GuidanceStore from .iteration_record import IterationRecord from .plugin import PLUGIN @@ -31,6 +38,7 @@ implement_prompt, perf_plan_prompt, plan_prompt, + repair_prompt, review_prompt, with_human_guidance, ) @@ -91,6 +99,7 @@ def __init__( self.champions = ChampionStore( cfg.state_dir / "champion", cfg.evaluator_bundle.spec.acceptance.noise_threshold, + cfg.evaluator_bundle.spec.benchmark_case_ids, ) def run(self) -> None: @@ -114,13 +123,15 @@ def run(self) -> None: self.store.update_run(current_iteration=0, current_phase="S_baseline") self.store.append_timeline("phase_start", {"iteration": 0, "phase": "S_baseline"}) self.baseline = self._ensure_triton_baseline() - self.champions.initialize_triton(self.baseline) + self.champions.initialize_triton(self.baseline["benchmark_report"]) self.initial_hip = self._ensure_initial_hip() - initial_score = dict(self.initial_hip.get("benchmark", {}).get("score") or {}) + initial_score = dict(self.initial_hip["benchmark"].get("score") or {}) if initial_score.get("passed"): promoted, reason, champion = self.champions.consider( - 0, self.cfg.state_dir / "certified" / "initial-hip" / "submission", - initial_score, + 0, + self.cfg.state_dir / "certified" / "initial-hip" / "submission", + self.initial_hip["benchmark_report"], + self.baseline["benchmark_report"], ) self.store.append_timeline("initial_hip_challenged", { "promoted": promoted, "reason": reason, "champion": champion, @@ -216,6 +227,27 @@ def _run_iteration(self, n: int) -> P.Outcome: ), ) + if not ( + compile_result.passed and correctness is not None and correctness.passed + ): + ok, repair_failure = self._agent_phase( + rec, "B_implement", role="repair", workdir=submission_dir, + prompt=repair_prompt(self.agent_req, submission_dir, n, test_feedback), + ) + if ok: + build_result, compile_result, correctness = self._test_phase( + rec, submission_dir, logs_dir + ) + test_feedback = self._write_feedback( + logs_dir, compile_result=compile_result, + correctness_result=correctness, + ) + rec.phases["C_test"]["repair_attempted"] = True + self._write(rec) + else: + rec.phases.setdefault("C_test", {})["repair_failure"] = repair_failure + self._write(rec) + if not compile_result.passed: return self._finish_failed( rec, @@ -230,17 +262,63 @@ def _run_iteration(self, n: int) -> P.Outcome: ) benchmark = self._evaluation_phase( - rec, "E_perf_test", "benchmark", submission_dir, + rec, "E_perf_test", submission_dir, build_result.artifact_dir, logs_dir, ) score = dict(benchmark.report.get("score") or {}) rec.score = score - rec.hardware_profile = dict(benchmark.report.get("hardware_profile") or {}) + if benchmark.infra_failure: + perf_feedback = self._write_feedback( + logs_dir, compile_result=compile_result, + correctness_result=correctness, benchmark_result=benchmark, + ) + del perf_feedback + return self._finish_failed( + rec, P.INFRA_FAIL, + benchmark.failure or "benchmark infrastructure failure", + ) + + diagnostic_ids = [] if benchmark.passed else list( + score.get("failed_case_ids") or [] + ) + + if self.profiler is not None and (benchmark.passed or diagnostic_ids): + diagnostic = self.profiler.run( + build_result.artifact_dir, logs_dir, role="candidate", + collection_mode="full", + case_ids=None if benchmark.passed else diagnostic_ids, + run_label="candidate-diagnostic", + implementation="candidate", + ) + if diagnostic.passed: + diagnostic_path = logs_dir / "candidate-diagnostic-hardware-profile.json" + rec.profile_report = make_report_reference( + self.cfg.state_dir, diagnostic_path + ) + benchmark.report["_profile_report"] = diagnostic.report + elif benchmark.passed: + benchmark.infra_failure = True + benchmark.passed = False + benchmark.failure = ( + diagnostic.failure or "promotable candidate full PMC archive failed" + ) + self._write_feedback( + logs_dir, compile_result=compile_result, + correctness_result=correctness, benchmark_result=benchmark, + ) + return self._finish_failed( + rec, P.INFRA_FAIL, benchmark.failure + ) promoted = False reason = benchmark.failure or "benchmark failed" champion = self.champions.load() if benchmark.passed: - promoted, reason, champion = self.champions.consider(n, submission_dir, score) + promoted, reason, champion = self.champions.consider( + n, + submission_dir, + rec.measurement_report, + rec.incumbent_measurement_report, + ) rec.promoted = promoted rec.champion_iteration = int(champion.get("iteration", 0)) self._write(rec) @@ -254,8 +332,6 @@ def _run_iteration(self, n: int) -> P.Outcome: ) self._perf_plan(rec, iter_dir, perf_feedback, promotion) - if benchmark.infra_failure: - return self._finish_failed(rec, P.INFRA_FAIL, benchmark.failure or "benchmark infrastructure failure") outcome = P.OK if promoted else P.PERF_REGRESSION return self._finish(rec, "success" if promoted else "not_promoted", outcome, reason if not promoted else None) @@ -329,60 +405,231 @@ def _evaluation_phase( self, rec: IterationRecord, phase: P.Phase, - evaluator_phase: str, submission_dir: Path, artifact_dir: Path, logs_dir: Path, ) -> EvaluationResult: self._start_phase(rec, phase) try: - result = self.evaluator.run( - evaluator_phase, - submission_dir, - artifact_dir, - logs_dir, - role="candidate", - build_fingerprint=self.builder.profile.fingerprint, - baseline_report=( - self.baseline.get("benchmark") if evaluator_phase == "benchmark" else None - ), + incumbent = self.champions.load() + if incumbent.get("kind") == "hip": + incumbent_build = self.builder.build( + self.champions.submission_dir, logs_dir / "incumbent-build" + ) + if not incumbent_build.passed: + raise RuntimeError( + incumbent_build.failure or "current Champion did not rebuild" + ) + incumbent_artifact = incumbent_build.artifact_dir + incumbent_impl = "candidate" + incumbent_fingerprint = self.builder.profile.fingerprint + else: + incumbent_artifact = self.cfg.state_dir / "baseline" / "runtime-artifacts" + incumbent_impl = "triton" + incumbent_fingerprint = "triton-jit" + incumbent_result, incumbent_ref, incumbent_profile_ref, _ = self._profile_benchmark( + incumbent_artifact, logs_dir, role="baseline", + build_fingerprint=incumbent_fingerprint, + report_label="incumbent", collection_mode="trace", + implementation=incumbent_impl, + ) + if not incumbent_result.passed or incumbent_ref is None: + raise RuntimeError( + incumbent_result.failure or "same-round Champion trace failed" + ) + rec.incumbent_measurement_report = incumbent_ref + if incumbent_profile_ref: + rec.incumbent_profile_report = incumbent_profile_ref + result, measurement_ref, profile_ref, profile_report = ( + self._profile_benchmark( + artifact_dir, + logs_dir, + role="candidate", + build_fingerprint=self.builder.profile.fingerprint, + baseline_report=incumbent_result.report, + collection_mode="trace", + ) ) - if ( - evaluator_phase == "benchmark" and result.passed - and self.profiler is not None + if result.passed and _near_promotion_boundary( + incumbent_result.report, result.report, + self.cfg.evaluator_bundle.spec.acceptance.noise_threshold, ): - profile_result = self.profiler.run( - artifact_dir, logs_dir, role="candidate" + incumbent_retry, _, _, _ = self._profile_benchmark( + incumbent_artifact, logs_dir, role="baseline", + build_fingerprint=incumbent_fingerprint, + report_label="incumbent-retest", collection_mode="trace", + implementation=incumbent_impl, ) - result.report["hardware_profile"] = profile_result.report - if not profile_result.passed and self.profiler.profile.required: - result = EvaluationResult( - evaluator_phase, False, result.report, - profile_result.failure or "required hardware profile failed", True, - ) + candidate_retry, _, _, _ = self._profile_benchmark( + artifact_dir, logs_dir, role="candidate", + build_fingerprint=self.builder.profile.fingerprint, + baseline_report=incumbent_retry.report, + report_label="candidate-retest", collection_mode="trace", + ) + if not incumbent_retry.passed or candidate_retry.infra_failure: + raise RuntimeError("boundary retest hipprof trace failed") + incumbent_combined = _combine_hipprof_reports( + incumbent_result.report, incumbent_retry.report + ) + candidate_combined = _combine_hipprof_reports( + result.report, candidate_retry.report + ) + incumbent_path = logs_dir / "incumbent-combined-benchmark-report.json" + candidate_path = logs_dir / "candidate-combined-benchmark-report.json" + write_json_atomic(incumbent_path, incumbent_combined) + incumbent_ref = make_report_reference(self.cfg.state_dir, incumbent_path) + rec.incumbent_measurement_report = incumbent_ref + result = self.evaluator.validate_benchmark_report( + candidate_combined, role="candidate", + build_fingerprint=self.builder.profile.fingerprint, + baseline_report=incumbent_combined, + ) + result.report["boundary_retested"] = True + write_json_atomic(candidate_path, result.report) + measurement_ref = make_report_reference(self.cfg.state_dir, candidate_path) + if measurement_ref: + rec.measurement_report = measurement_ref + if profile_ref: + rec.profile_report = profile_ref + if profile_report: + result.report["_profile_report"] = profile_report except Exception as exc: # noqa: BLE001 result = EvaluationResult( - evaluator_phase, False, {}, f"evaluator crashed: {exc!r}", True + "benchmark", False, {}, f"hipprof evaluation crashed: {exc!r}", True ) outcome = P.OK if result.passed else ( - P.INFRA_FAIL if result.infra_failure else ( - P.PERF_REGRESSION if evaluator_phase == "benchmark" else P.LOGIC_FAIL - ) + P.INFRA_FAIL if result.infra_failure else P.PERF_REGRESSION ) summary: Dict[str, Any] = { - "report": str(logs_dir / f"candidate-{evaluator_phase}-report.json"), + "report": str(logs_dir / "candidate-benchmark-report.json"), "build_fingerprint": self.builder.profile.fingerprint, + "score": result.report.get("score"), + "measurement_report": dict(rec.measurement_report), + "profile_report": dict(rec.profile_report), + "incumbent_measurement_report": dict(rec.incumbent_measurement_report), } - if evaluator_phase == "benchmark": - summary["score"] = result.report.get("score") - summary["hardware_profile"] = str( - logs_dir / "candidate-hardware-profile.json" - ) - if evaluator_phase == "correctness": - summary["summary"] = result.report.get("summary") self._end_phase(rec, phase, outcome, result.failure, summary) return result + def _profile_benchmark( + self, + artifact_dir: Path, + report_dir: Path, + *, + role: str, + build_fingerprint: str, + baseline_report: Optional[Dict[str, Any]] = None, + report_label: Optional[str] = None, + collection_mode: str = "trace", + implementation: Optional[str] = None, + ) -> tuple[ + EvaluationResult, + Optional[ReportReference], + Optional[ReportReference], + Dict[str, Any], + ]: + report_dir.mkdir(parents=True, exist_ok=True) + label = report_label or role + benchmark_path = report_dir / f"{label}-benchmark-report.json" + if self.profiler is None: + result = EvaluationResult( + "benchmark", False, {}, "required hipprof profiler is unavailable", True + ) + return result, None, None, {} + + profiled = self.profiler.run( + artifact_dir, report_dir, role=role, + collection_mode=collection_mode, run_label=label, + implementation=implementation, + ) + + profile_report = dict(profiled.report) + profile_path = report_dir / f"{label}-hardware-profile.json" + try: + profile_ref = make_report_reference(self.cfg.state_dir, profile_path) + except RuntimeError as exc: + result = EvaluationResult( + "benchmark", False, {}, f"hipprof report is unavailable: {exc}", True + ) + return result, None, None, profile_report + + report: Dict[str, Any] = { + "schema_version": 2, + "passed": bool(profiled.passed), + "methodology": dict(self.cfg.evaluator_bundle.spec.benchmark_protocol), + "timing_source": "hipprof GPU kernel DurationNs", + "timed_scope": "operator_gpu_dispatches_only", + "profile_report": dict(profile_ref), + "cases": [], + } + if not profiled.passed: + write_json_atomic(benchmark_path, report) + measurement_ref = make_report_reference(self.cfg.state_dir, benchmark_path) + result = EvaluationResult( + "benchmark", + False, + report, + profiled.failure or "required hipprof profile failed", + True, + ) + return result, measurement_ref, profile_ref, profile_report + + specs = { + spec.id: spec for spec in self.cfg.evaluator_bundle.spec.benchmark_cases + } + timing_cases = profile_report.get("timing_cases") or [] + ids = [ + str(case.get("id") or "") + for case in timing_cases + if isinstance(case, dict) + ] + duplicates = sorted({case_id for case_id in ids if ids.count(case_id) > 1}) + expected = self.cfg.evaluator_bundle.spec.benchmark_case_ids + missing = sorted(set(expected) - set(ids)) + unexpected = sorted(set(ids) - set(expected)) + invalid = len(ids) != len(timing_cases) or "" in ids + if missing or unexpected or duplicates or invalid: + report["passed"] = False + report["profile_case_errors"] = { + "missing": missing, + "unexpected": unexpected, + "duplicate": duplicates, + "invalid": invalid, + } + write_json_atomic(benchmark_path, report) + measurement_ref = make_report_reference(self.cfg.state_dir, benchmark_path) + failure = ( + "hipprof timing cases are incomplete: " + f"missing={missing}, unexpected={unexpected}, " + f"duplicate={duplicates}, invalid={invalid}" + ) + result = EvaluationResult("benchmark", False, report, failure, True) + return result, measurement_ref, profile_ref, profile_report + + cases: List[Dict[str, Any]] = [] + for raw in timing_cases: + case_id = str(raw["id"]) + spec = specs[case_id] + case = dict(raw) + case["timing_source"] = "hipprof GPU kernel DurationNs" + if spec.shape is not None: + case["shape"] = dict(spec.shape) + if spec.flops is not None: + case["flops"] = spec.flops + if spec.bytes is not None: + case["bytes"] = spec.bytes + cases.append(case) + report["cases"] = cases + result = self.evaluator.validate_benchmark_report( + report, + role=role, + build_fingerprint=build_fingerprint, + baseline_report=baseline_report, + ) + write_json_atomic(benchmark_path, result.report) + measurement_ref = make_report_reference(self.cfg.state_dir, benchmark_path) + return result, measurement_ref, profile_ref, profile_report + def _test_phase( self, rec: IterationRecord, submission_dir: Path, logs_dir: Path, ) -> tuple[BuildResult, EvaluationResult, Optional[EvaluationResult]]: @@ -433,27 +680,56 @@ def _ensure_triton_baseline(self) -> Dict[str, Any]: artifact_dir = baseline_dir / "runtime-artifacts" manifest_path = baseline_dir / "baseline-manifest.json" if manifest_path.is_file(): - manifest = json.loads(manifest_path.read_text(encoding="utf-8")) - digest = manifest.pop("manifest_sha256", None) - actual = _canonical_digest(manifest) - if digest != actual: + stored = json.loads(manifest_path.read_text(encoding="utf-8")) + digest = stored.get("manifest_sha256") + manifest = { + key: value for key, value in stored.items() + if key != "manifest_sha256" + } + if digest != _canonical_digest(manifest): raise RuntimeError("frozen baseline manifest changed") if manifest.get("evaluator_digest") != self.cfg.evaluator_bundle.digest: raise RuntimeError("baseline evaluator differs from active evaluator") - if self.profiler is not None: - expected_profile = self.profiler.profile.fingerprint - actual_profile = ( - manifest.get("hardware_profile") or {} - ).get("profile_fingerprint") - if actual_profile != expected_profile: - raise RuntimeError("baseline profiler differs from active hardware profile") if manifest.get("implementation") != "triton": raise RuntimeError("baseline is not the certified Triton implementation") + benchmark_ref = manifest.get("benchmark_report") + profile_ref = manifest.get("profile_report") + if not isinstance(benchmark_ref, dict): + benchmark_ref = make_report_reference( + self.cfg.state_dir, + baseline_dir / "baseline-benchmark-report.json", + ) + if not isinstance(profile_ref, dict): + profile_ref = make_report_reference( + self.cfg.state_dir, + baseline_dir / "baseline-hardware-profile.json", + ) + benchmark = load_report_reference(self.cfg.state_dir, benchmark_ref) + profile = load_report_reference(self.cfg.state_dir, profile_ref) + validated = self.evaluator.validate_benchmark_report( + benchmark, + role="baseline", + build_fingerprint="triton-jit", + ) + if not validated.passed: + raise RuntimeError(validated.failure or "frozen Triton benchmark is invalid") + if self.profiler is None: + raise RuntimeError("required hipprof profiler is unavailable") + if profile.get("profile_fingerprint") != self.profiler.profile.fingerprint: + raise RuntimeError( + "baseline profiler differs from active hardware profile" + ) self.store.append_timeline( - "baseline_reused", - {"implementation": "triton"}, + "baseline_reused", {"implementation": "triton"} ) - return manifest + return { + **manifest, + "manifest_sha256": digest, + "benchmark_report": dict(benchmark_ref), + "profile_report": dict(profile_ref), + "benchmark": validated.report, + "hardware_profile": profile, + } baseline_dir.mkdir(parents=True, exist_ok=True) submission.mkdir(parents=True, exist_ok=True) @@ -468,40 +744,42 @@ def _ensure_triton_baseline(self) -> Dict[str, Any]: ) if not correctness.passed: raise RuntimeError(correctness.failure or "Triton baseline failed correctness") - benchmark = self.evaluator.run( - "benchmark", submission, artifact_dir, baseline_dir, - role="baseline", build_fingerprint="triton-jit", + benchmark, benchmark_ref, profile_ref, profile = self._profile_benchmark( + artifact_dir, + baseline_dir, + role="baseline", + build_fingerprint="triton-jit", + collection_mode="full", ) - if not benchmark.passed: - raise RuntimeError(benchmark.failure or "Triton baseline benchmark failed") - hardware_profile: Dict[str, Any] = {} - if self.profiler is not None: - profiled = self.profiler.run( - artifact_dir, baseline_dir, role="baseline" + if not benchmark.passed or benchmark_ref is None or profile_ref is None: + raise RuntimeError( + benchmark.failure or "Triton baseline hipprof benchmark failed" ) - hardware_profile = profiled.report - if not profiled.passed and self.profiler.profile.required: - raise RuntimeError(profiled.failure or "Triton hardware profile failed") payload = { - "schema_version": 2, + "schema_version": 3, "implementation": "triton", "certified_at": time.time(), "build_fingerprint": "triton-jit", "evaluator_digest": self.cfg.evaluator_bundle.digest, "correctness": correctness.report, - "benchmark": benchmark.report, - "hardware_profile": hardware_profile, + "benchmark_report": dict(benchmark_ref), + "profile_report": dict(profile_ref), } payload["manifest_sha256"] = _canonical_digest(payload) - manifest_path.write_text(json.dumps(payload, indent=2), encoding="utf-8") + write_json_atomic(manifest_path, payload) self.store.append_timeline( "baseline_certified", { "implementation": "triton", "benchmark_cases": len(benchmark.report.get("cases") or []), + "measurement_report": benchmark_ref, }, ) - return payload + return { + **payload, + "benchmark": benchmark.report, + "hardware_profile": profile, + } def _ensure_initial_hip(self) -> Dict[str, Any]: """Independently certify the user-provided HIP optimization seed.""" @@ -509,13 +787,54 @@ def _ensure_initial_hip(self) -> Dict[str, Any]: submission = root / "submission" manifest_path = root / "initial-hip-manifest.json" if manifest_path.is_file(): - manifest = json.loads(manifest_path.read_text(encoding="utf-8")) - digest = manifest.pop("manifest_sha256", None) + stored = json.loads(manifest_path.read_text(encoding="utf-8")) + digest = stored.get("manifest_sha256") + manifest = { + key: value for key, value in stored.items() + if key != "manifest_sha256" + } if digest != _canonical_digest(manifest): raise RuntimeError("frozen Initial HIP manifest changed") if manifest.get("submission_digest") != _tree_digest(submission): raise RuntimeError("frozen Initial HIP submission changed") - return manifest + if manifest.get("evaluator_digest") != self.cfg.evaluator_bundle.digest: + raise RuntimeError("Initial HIP evaluator differs from active evaluator") + benchmark_ref = manifest.get("benchmark_report") + profile_ref = manifest.get("profile_report") + if not isinstance(benchmark_ref, dict): + benchmark_ref = make_report_reference( + self.cfg.state_dir, + root / "candidate-benchmark-report.json", + ) + if not isinstance(profile_ref, dict): + profile_ref = make_report_reference( + self.cfg.state_dir, + root / "candidate-hardware-profile.json", + ) + benchmark = load_report_reference(self.cfg.state_dir, benchmark_ref) + profile = load_report_reference(self.cfg.state_dir, profile_ref) + validated = self.evaluator.validate_benchmark_report( + benchmark, + role="candidate", + build_fingerprint=self.builder.profile.fingerprint, + baseline_report=self.baseline["benchmark"], + ) + if validated.infra_failure: + raise RuntimeError(validated.failure or "Initial HIP benchmark is invalid") + if self.profiler is None: + raise RuntimeError("required hipprof profiler is unavailable") + if profile.get("profile_fingerprint") != self.profiler.profile.fingerprint: + raise RuntimeError( + "Initial HIP profiler differs from active hardware profile" + ) + return { + **manifest, + "manifest_sha256": digest, + "benchmark_report": dict(benchmark_ref), + "profile_report": dict(profile_ref), + "benchmark": validated.report, + "hardware_profile": profile, + } initial = self.cfg.initial_submission if initial is None or not initial.is_dir(): @@ -531,19 +850,20 @@ def _ensure_initial_hip(self) -> Dict[str, Any]: ) if not correctness.passed: raise RuntimeError(correctness.failure or "Initial HIP failed correctness") - benchmark = self.evaluator.run( - "benchmark", submission, build_result.artifact_dir, root, - role="candidate", build_fingerprint=self.builder.profile.fingerprint, + benchmark, benchmark_ref, profile_ref, profile = self._profile_benchmark( + build_result.artifact_dir, + root, + role="candidate", + build_fingerprint=self.builder.profile.fingerprint, baseline_report=self.baseline["benchmark"], + collection_mode="full", ) - hardware_profile: Dict[str, Any] = {} - if self.profiler is not None: - profiled = self.profiler.run(build_result.artifact_dir, root, role="candidate") - hardware_profile = profiled.report - if not profiled.passed and self.profiler.profile.required: - raise RuntimeError(profiled.failure or "Initial HIP hardware profile failed") + if benchmark.infra_failure or benchmark_ref is None or profile_ref is None: + raise RuntimeError( + benchmark.failure or "Initial HIP hipprof benchmark failed" + ) payload = { - "schema_version": 1, + "schema_version": 2, "implementation": "initial-hip", "certified_at": time.time(), "build_fingerprint": self.builder.profile.fingerprint, @@ -551,16 +871,21 @@ def _ensure_initial_hip(self) -> Dict[str, Any]: "submission_digest": _tree_digest(submission), "compile": build_result.report, "correctness": correctness.report, - "benchmark": benchmark.report, - "hardware_profile": hardware_profile, + "benchmark_report": dict(benchmark_ref), + "profile_report": dict(profile_ref), } payload["manifest_sha256"] = _canonical_digest(payload) - manifest_path.write_text(json.dumps(payload, indent=2), encoding="utf-8") + write_json_atomic(manifest_path, payload) self.store.append_timeline("initial_hip_certified", { "benchmark_cases": len(benchmark.report.get("cases") or []), "passed_gates": bool((benchmark.report.get("score") or {}).get("passed")), + "measurement_report": benchmark_ref, }) - return payload + return { + **payload, + "benchmark": benchmark.report, + "hardware_profile": profile, + } def _review( self, @@ -681,7 +1006,7 @@ def _write_feedback( "score": score, "methodology": benchmark_result.report.get("methodology") or {}, "hardware_profile": _agent_profile_feedback( - benchmark_result.report.get("hardware_profile") or {} + benchmark_result.report.get("_profile_report") or {} ), } if promotion: @@ -778,7 +1103,11 @@ def _agent_profile_feedback(report: Dict[str, Any]) -> Dict[str, Any]: "gpu_arch": report.get("gpu_arch"), "tool": report.get("tool"), "counter_groups": report.get("counter_groups") or [], - "cases": report.get("cases") or [], + "cases": [ + {key: value for key, value in case.items() + if key != "operator_samples_us"} + for case in report.get("cases") or [] if isinstance(case, dict) + ], } @@ -790,6 +1119,70 @@ def _canonical_digest(data: Dict[str, Any]) -> str: ).hexdigest() +def _near_promotion_boundary( + incumbent: Dict[str, Any], candidate: Dict[str, Any], threshold: float, +) -> bool: + incumbent_by_id = { + str(case.get("id")): case for case in incumbent.get("cases") or [] + if isinstance(case, dict) and case.get("id") + } + for case in candidate.get("cases") or []: + case_id = str(case.get("id") or "") + if case_id not in incumbent_by_id: + continue + old = float(incumbent_by_id[case_id]["latency_ms"]) + new = float(case["latency_ms"]) + improvement = 1.0 - new / old + if abs(improvement - threshold) <= threshold: + return True + return False + + +def _combine_hipprof_reports( + first: Dict[str, Any], second: Dict[str, Any], +) -> Dict[str, Any]: + """Combine equal-sized hipprof batches without shape weighting.""" + other = { + str(case.get("id")): case for case in second.get("cases") or [] + if isinstance(case, dict) and case.get("id") + } + combined_cases = [] + for raw in first.get("cases") or []: + case = dict(raw) + peer = other.get(str(case.get("id") or "")) + if peer is None: + raise RuntimeError("boundary retest cases differ") + samples = [ + float(value) for value in case.get("operator_samples_ms") or [] + ] + [ + float(value) for value in peer.get("operator_samples_ms") or [] + ] + if not samples: + samples = [float(case["latency_ms"]), float(peer["latency_ms"])] + mean = statistics.fmean(samples) + case.update({ + "latency_ms": mean, + "latency_mean_ms": mean, + "latency_median_ms": statistics.median(samples), + "latency_stddev_ms": statistics.stdev(samples) if len(samples) > 1 else 0.0, + "latency_cv": ( + statistics.stdev(samples) / mean if len(samples) > 1 and mean > 0 else 0.0 + ), + "latency_min_ms": min(samples), + "latency_max_ms": max(samples), + "operator_samples_ms": samples, + "sample_count": len(samples), + "measurement_batches": 2, + }) + combined_cases.append(case) + result = dict(first) + result["cases"] = combined_cases + result["measurement_batches"] = 2 + result["aggregation"] = "arithmetic mean of all raw hipprof DurationNs operator samples" + result.pop("score", None) + return result + + def _tree_digest(root: Path) -> str: digest = hashlib.sha256() for path in sorted(root.rglob("*")): diff --git a/metainfer/tasks/opt_GEMM_kernel/orchestrator/profiler.py b/metainfer/tasks/opt_GEMM_kernel/orchestrator/profiler.py index c4569357..149b6ce8 100644 --- a/metainfer/tasks/opt_GEMM_kernel/orchestrator/profiler.py +++ b/metainfer/tasks/opt_GEMM_kernel/orchestrator/profiler.py @@ -11,6 +11,7 @@ import resource import shutil import subprocess +import sys from dataclasses import asdict, dataclass from pathlib import Path from typing import Any, Dict, Iterable, List, Mapping, Optional @@ -25,6 +26,22 @@ def _disable_core_dump() -> None: resource.setrlimit(resource.RLIMIT_CORE, (0, 0)) +def _python_executable() -> str: + """Return the real interpreter even when a launcher rewrites argv[0].""" + candidate = Path(sys.executable) if sys.executable else None + if candidate is not None and candidate.name.lower().startswith("python"): + return str(candidate.resolve()) + try: + executable = Path("/proc/self/exe").resolve(strict=True) + except OSError as exc: + raise ProfilerError("cannot resolve the Python executable") from exc + if not executable.name.lower().startswith("python"): + raise ProfilerError( + f"resolved process executable is not Python: {executable}" + ) + return str(executable) + + def _version(executable: str) -> str: # rocprofv3 and newer rocprof releases expose a conventional version # option. The legacy RPL rocprof shipped by DTK only exposes ``-h``; @@ -134,16 +151,14 @@ def resolve( if executable is None: if raw.get("required", True): raise ProfilerError( - "Hygon K100/gfx928 requires hipprof, rocprofv3, or rocprof " - "on the target node" + "Hygon K100/gfx928 requires hipprof on the target node" ) return None - if "hipprof" in executable.name: - kind = "hipprof" - elif "rocprofv3" in executable.name: - kind = "rocprofv3" - else: - kind = "rocprof" + if "hipprof" not in executable.name: + raise ProfilerError( + f"Hygon K100/gfx928 timing requires hipprof, got {executable}" + ) + kind = "hipprof" available = _available_counters(executable, kind) configured = [list(map(str, group)) for group in raw.get("counter_groups") or []] groups = ( @@ -180,6 +195,10 @@ def resolve( return cls(**data) def verify(self) -> None: + if self.tool_kind != "hipprof": + raise ProfilerError( + f"Hygon K100/gfx928 timing requires hipprof, got {self.tool_kind}" + ) data = asdict(self) expected = _fingerprint({**data, "fingerprint": ""}) if expected != self.fingerprint: @@ -203,10 +222,12 @@ def __init__( *, private_env: Mapping[str, str], harness_argv: Optional[List[str]] = None, + benchmark_protocol: Optional[Mapping[str, Any]] = None, ) -> None: self.profile = profile self.private_env = dict(private_env) self.harness_argv = harness_argv + self.benchmark_protocol = dict(benchmark_protocol or {}) def run( self, @@ -214,22 +235,48 @@ def run( output_dir: Path, *, role: str, + collection_mode: str = "trace", + case_ids: Optional[List[str]] = None, + run_label: Optional[str] = None, + implementation: Optional[str] = None, ) -> ProfileResult: try: self.profile.verify() except Exception as exc: # noqa: BLE001 return ProfileResult(False, {}, str(exc)) - if self.harness_argv is not None: - harness_cmd = list(self.harness_argv) - else: - harness = artifact_dir / "metainfer_gemm_harness" - if not harness.is_file(): - return ProfileResult(False, {}, f"native harness is missing: {harness}") - harness_cmd = [str(harness.resolve())] + if self.profile.tool_kind != "hipprof": + return ProfileResult( + False, {}, + f"K100 performance evaluation requires hipprof, got {self.profile.tool_kind}", + ) + if self.harness_argv is None or len(self.harness_argv) < 2: + return ProfileResult(False, {}, "frozen hipprof workload driver is missing") + harness_cmd = list(self.harness_argv) - root = output_dir / f"{role}-hardware-profile" + if collection_mode not in {"trace", "full"}: + return ProfileResult(False, {}, f"invalid collection mode: {collection_mode}") + root = output_dir / (run_label or f"{role}-hardware-profile") root.mkdir(parents=True, exist_ok=True) + evaluator = Path(harness_cmd[1]).resolve() + suite = evaluator.with_name("run_hipprof_suite.py") + analyzer = evaluator.with_name("analyze_hipprof_suite.py") + if not suite.is_file() or not analyzer.is_file(): + return ProfileResult(False, {}, "task-local hipprof suite or analyzer is missing") + return self._run_hipprof_suite( + artifact_dir, output_dir, root, role, suite, analyzer, + collection_mode=collection_mode, case_ids=case_ids, + report_label=run_label or role, implementation=implementation) + + def _run_legacy_profile_route( + self, + artifact_dir: Path, + output_dir: Path, + root: Path, + role: str, + harness_cmd: List[str], + ) -> ProfileResult: + """Retained only for parsing historical profiler fixtures.""" cases: List[Dict[str, Any]] = [] commands: List[List[str]] = [] for case_id in self.profile.representative_cases: @@ -302,6 +349,175 @@ def run( _write_json(output_dir / f"{role}-hardware-profile.json", report) return ProfileResult(True, report) + def _run_hipprof_suite( + self, + artifact_dir: Path, + output_dir: Path, + root: Path, + role: str, + suite: Path, + analyzer: Path, + *, + collection_mode: str, + case_ids: Optional[List[str]], + report_label: str, + implementation: Optional[str], + ) -> ProfileResult: + """Run task-local full trace + PMC/read/write collection. + + The frozen harness owns case enumeration and analysis. The system still + owns the resolved hipprof executable, candidate artifact, weights, and + output directory; no external benchmark checkout is consulted. + """ + implementation = implementation or ( + "triton" if role == "baseline" else "candidate" + ) + if implementation not in {"triton", "candidate"}: + return ProfileResult(False, {}, f"invalid implementation: {implementation}") + python_executable = _python_executable() + command = [ + python_executable, str(suite), + "--hipprof", self.profile.executable, + "--output-dir", str(root.resolve()), + "--implementations", implementation, + "--passes", collection_mode, + ] + if case_ids: + command.extend(["--case-ids", ",".join(case_ids)]) + env = dict(os.environ) + env.update(self.private_env) + env.update({ + "METAINFER_EVALUATION_PHASE": "profile-batch", + "METAINFER_EVALUATION_ROLE": role, + "METAINFER_BUILD_ARTIFACT_DIR": str(artifact_dir.resolve()), + "METAINFER_REPORT_PATH": str((root / "suite-report.json").resolve()), + "METAINFER_BENCHMARK_PROTOCOL": json.dumps( + self.benchmark_protocol, sort_keys=True + ), + "PYTHONDONTWRITEBYTECODE": "1", + }) + try: + proc = subprocess.run( + command, cwd=str(suite.parent), env=env, text=True, + stdout=subprocess.PIPE, stderr=subprocess.PIPE, + timeout=2400 if collection_mode == "full" else 900, + check=False, preexec_fn=_disable_core_dump, + ) + except subprocess.TimeoutExpired as exc: + report = self._report([], [command]) + _write_json(output_dir / f"{report_label}-hardware-profile.json", report) + return ProfileResult(False, report, f"hipprof suite timed out: {exc}") + (root / "suite.stdout.log").write_text(proc.stdout or "", encoding="utf-8") + (root / "suite.stderr.log").write_text(proc.stderr or "", encoding="utf-8") + if proc.returncode: + report = self._report([], [command]) + _write_json(output_dir / f"{report_label}-hardware-profile.json", report) + return ProfileResult(False, report, "task-local hipprof suite failed") + + analyze_command = [python_executable, str(analyzer), str(root.resolve())] + analyzed = subprocess.run( + analyze_command, cwd=str(analyzer.parent), env=env, text=True, + stdout=subprocess.PIPE, stderr=subprocess.PIPE, + timeout=300, check=False, preexec_fn=_disable_core_dump, + ) + (root / "analyze.stdout.log").write_text( + analyzed.stdout or "", encoding="utf-8") + (root / "analyze.stderr.log").write_text( + analyzed.stderr or "", encoding="utf-8") + if analyzed.returncode: + report = self._report([], [command, analyze_command]) + _write_json(output_dir / f"{report_label}-hardware-profile.json", report) + return ProfileResult(False, report, "task-local hipprof analysis failed") + + metrics = json.loads((root / "metrics.json").read_text(encoding="utf-8")) + cases = [] + for row in metrics.get("rows", []): + read_bytes = float(row.get("hbm_read_bytes") or 0) + write_bytes = float(row.get("hbm_write_bytes") or 0) + try: + kernel_breakdown = json.loads( + row.get("trace_kernel_means_json") or "{}" + ) + except (TypeError, json.JSONDecodeError): + kernel_breakdown = {} + cases.append({ + "id": row["case_id"], + "duration_ns": float(row["operator_mean_us"]) * 1000.0, + "latency_mean_us": float(row["operator_mean_us"]), + "latency_median_us": float(row.get("operator_median_us") or 0), + "latency_stddev_us": float(row.get("operator_stddev_us") or 0), + "latency_cv": float(row.get("operator_cv") or 0), + "latency_min_us": float(row.get("operator_min_us") or 0), + "latency_max_us": float(row.get("operator_max_us") or 0), + "operator_samples_us": list(row.get("operator_samples_us") or []), + "dispatch_count": int(row.get("trace_dispatches_per_call") or 0), + "pmc_dispatch_count": int(row.get("pmc_dispatch_count") or 0), + "kernel_name": row.get("main_kernel", ""), + "kernel_breakdown_us": kernel_breakdown, + "vgpr_count": int(row.get("vgpr") or 0), + "agpr_count": int(row.get("agpr") or 0), + "sgpr_count": int(row.get("sgpr") or 0), + "lds_bytes": int(row.get("lds_bytes") or 0), + "scratch_bytes": int(row.get("scratch_bytes") or 0), + "grid_size": int(row.get("grid_size") or 0), + "workgroup_size": int(row.get("workgroup_size") or 0), + "wave_size": int(row.get("wave_size") or 0), + "waves_per_workgroup": row.get("waves_per_workgroup"), + "occupancy_pct": row.get("occupancy_pct"), + "l2_hit_pct": row.get("l2_hit_pct"), + "hbm_read_bytes": read_bytes, + "hbm_write_bytes": write_bytes, + "hbm_read_gbps": row.get("hbm_read_gbs"), + "hbm_write_gbps": row.get("hbm_write_gbs"), + "measured_bandwidth_gbps": row.get("hbm_total_gbs"), + "compute_busy_pct": None, + "matrix_instructions": None, + "valu_instructions": None, + "counters": { + "HBM_READ_BYTES": read_bytes, + "HBM_WRITE_BYTES": write_bytes, + }, + }) + report = self._report(cases, [command, analyze_command]) + # The task-local analyzer has already validated the exact requested + # case set. Representative-case coverage applies only to the legacy + # route, not to scoped diagnostic collections. + report["passed"] = bool(cases) + report["collection_mode"] = collection_mode + report["collection_method"] = ( + "task-local hipprof trace plus pmc/pmc-read/pmc-write" + if collection_mode == "full" else "task-local hipprof trace" + ) + report["pass_records"] = metrics.get("pass_records") or [] + report["full_metrics_csv"] = str((root / "metrics.csv").resolve()) + report["timing_source"] = "hipprof kernel DurationNs" + report["timing_cases"] = [ + { + "id": row["case_id"], + "latency_ms": float(row["operator_mean_us"]) / 1000.0, + "latency_mean_ms": float(row["operator_mean_us"]) / 1000.0, + "latency_median_ms": float(row.get("operator_median_us") or 0) / 1000.0, + "latency_stddev_ms": float(row.get("operator_stddev_us") or 0) / 1000.0, + "latency_cv": float(row.get("operator_cv") or 0), + "latency_min_ms": float(row.get("operator_min_us") or 0) / 1000.0, + "latency_max_ms": float(row.get("operator_max_us") or 0) / 1000.0, + "operator_samples_ms": [ + float(value) / 1000.0 + for value in row.get("operator_samples_us") or [] + ], + "sample_count": len(row.get("operator_samples_us") or []), + "kernel_name": row.get("main_kernel") or row.get("timed_kernel", ""), + "dispatch_count": int(row.get("trace_dispatches_per_call") or 0), + "kernel_breakdown_us": json.loads( + row.get("trace_kernel_means_json") or "{}" + ), + } + for row in metrics.get("rows", []) + ] + _write_json(output_dir / f"{report_label}-hardware-profile.json", report) + return ProfileResult(report["passed"], report, + None if report["passed"] else "profile cases missing") + def _command( self, harness_cmd: List[str], case_id: str, counters: List[str], output: Path, ) -> List[str]: @@ -332,9 +548,13 @@ def _command( ] def _report(self, cases: List[Dict[str, Any]], commands: List[List[str]]) -> Dict[str, Any]: + observed = { + str(case.get("id") or case.get("case_id") or "") for case in cases + } + expected = set(self.profile.representative_cases) return { "schema_version": 1, - "passed": len(cases) == len(self.profile.representative_cases), + "passed": bool(cases) and expected.issubset(observed), "profile_id": self.profile.id, "label": self.profile.label, "gpu_arch": self.profile.gpu_arch, diff --git a/metainfer/tasks/opt_GEMM_kernel/orchestrator/prompts.py b/metainfer/tasks/opt_GEMM_kernel/orchestrator/prompts.py index 9ef3898e..f413e493 100644 --- a/metainfer/tasks/opt_GEMM_kernel/orchestrator/prompts.py +++ b/metainfer/tasks/opt_GEMM_kernel/orchestrator/prompts.py @@ -54,12 +54,23 @@ def plan_prompt( {feedback} ``` -Read the contract and evaluation-protocol notebooks, then inspect the current -`submission/`. If `perf_plan.md` exists, it is the previous iteration's -F-phase recommendation: evaluate it against the current champion and public -feedback instead of ignoring it. Write the new `plan.md` in the iteration -directory. Choose one bounded, testable change, state the expected affected -shapes, risks, and rollback rule. Do not edit `submission/` in this phase. +Inspect the current `submission/`, the frozen public shapes, and the latest +per-shape evidence before reading optimization notes. Treat notebooks as +historical evidence, known constraints, and candidate hypotheses—not as an +exhaustive search space, a required dispatch recipe, or a substitute for +analyzing the current source and hardware. Notebook absolute timings are not +cross-machine service-level targets. If `perf_plan.md` exists, evaluate its +hypothesis against the current Champion, operator latency, dispatch breakdown, +and PMC/resource evidence rather than applying it mechanically. + +Write the new `plan.md` in the iteration directory. Select one evidence-backed, +bounded, measurable, reversible change; novelty is not a goal. State the exact +affected shape IDs, numerical latency expectation, source-level mechanism, +expected counter/resource movement, named control shapes expected to remain +unchanged, risks, and rollback rule. Valid directions include reducing summed GPU operator +time, reducing dispatch/reduction work, improving HBM or L2 efficiency, +reducing VGPR/AGPR/SGPR/LDS/scratch pressure, or improving parallelism when the +available evidence supports it. Do not edit `submission/` in this phase. {BOUNDARY} """ @@ -146,16 +157,53 @@ def perf_plan_prompt( {json.dumps(promotion, indent=2, ensure_ascii=False)} ``` -Write `perf_plan.md` in the iteration directory. Identify the shapes that -improved or regressed, distinguish measurement noise from a plausible kernel -bottleneck, use the system-provided hardware counters when present, and -propose one bounded next optimization. Do not edit +Write `perf_plan.md` in the iteration directory. Start from the current source +and every public shape's hipprof GPU operator time. For multi-dispatch calls, +interpret total operator latency and the per-kernel contribution breakdown; +do not treat the longest kernel or PMC replay duration as latency. Correlate +regressions with HBM read/write traffic and bandwidth, L2 behavior, +VGPR/AGPR/SGPR, LDS, scratch, dispatch count, and occupancy/wave information +only when the profiler actually reports it. + +Identify every failed shape and use reported dispersion/CV and any boundary +retest to distinguish noise from a plausible bottleneck. Rank up to three +evidence-backed hypotheses, then recommend one bounded, measurable, reversible +next optimization. Novelty is not a goal: prefer the strongest measured +evidence even when the direction is already documented. Include numerical +latency expectations for affected shapes, named control shapes expected to stay +unchanged, expected counter/resource changes, and an explicit rollback rule. Do not edit `submission/`; the next A/B phases execute the new plan. {BOUNDARY} """ +def repair_prompt( + req: Dict[str, Any], submission_dir: Path, iteration: int, + feedback: Dict[str, Any], +) -> str: + return f"""You are making the single allowed build/correctness repair for GEMM iteration {iteration}. + +Writable submission directory: {submission_dir} +Sanitized compiler/correctness evidence: +```json +{json.dumps(feedback, indent=2, ensure_ascii=False)} +``` + +Make only the smallest source change needed to address the evidenced failure. +Do not introduce a new optimization, broaden the original plan, change the ABI, +or modify any evaluator/profiler file. Update `CHANGELOG.md` with the repair. +If the evidence is insufficient, leave source unchanged and explain why there. + +Public requirements: +```json +{_requirements(req)} +``` + +{BOUNDARY} +""" + + def with_human_guidance(prompt: str, items: List[Dict[str, Any]]) -> str: """Append live user steering while reasserting the evaluator boundary.""" if not items: @@ -170,9 +218,11 @@ def with_human_guidance(prompt: str, items: List[Dict[str, Any]]) -> str: The task owner submitted the following optimization ideas while the task was running. Treat them as high-priority hypotheses within the public GEMM -contract. Inspect the current code and evidence before applying them. If an -idea is unsafe, incompatible with the ABI, or contradicted by measurements, -explain that in the plan/changelog instead of silently forcing it. +contract, not as instructions to bypass independent analysis. Inspect the +current code, target shapes, hipprof operator timing, dispatch breakdown, PMC +evidence, and hardware constraints before applying them. If an idea is unsafe, +incompatible with the ABI, or contradicted by measurements, explain that in +the plan/changelog instead of silently forcing it. {rendered} diff --git a/metainfer/tasks/opt_GEMM_kernel/server/_state_readers.py b/metainfer/tasks/opt_GEMM_kernel/server/_state_readers.py index ec022d3c..2b918bf8 100644 --- a/metainfer/tasks/opt_GEMM_kernel/server/_state_readers.py +++ b/metainfer/tasks/opt_GEMM_kernel/server/_state_readers.py @@ -1,15 +1,21 @@ -"""Read the task-owned iteration, score and champion schemas.""" +"""Read task-owned reports and derive public per-shape views.""" from __future__ import annotations -import json import copy +import json +import math from pathlib import Path from typing import Any, Dict, List, Optional +from metainfer.orchestrator.requirements import req_field + from ..orchestrator import phases +from ..orchestrator.evaluator.champion import ( + champion_report_reference, + load_report_reference, +) from ..orchestrator.evaluator.spec import BenchmarkCaseSpec, KernelTaskSpec, SpecError -from metainfer.orchestrator.requirements import req_field def _json(path: Path, default: Any) -> Any: @@ -21,21 +27,44 @@ def _json(path: Path, default: Any) -> Any: def read_iterations(state_dir: Path) -> List[Dict[str, Any]]: spec = _task_spec(state_dir) + baseline = _baseline_report(state_dir) records = [ value for path in sorted((state_dir / "iterations").glob("*.json")) if isinstance((value := _json(path, None)), dict) ] if (state_dir / "iterations").is_dir() else [] - return [_public_record(record, spec) for record in records] + return [ + _public_record(state_dir, record, spec, baseline) + for record in records + ] def read_iteration(state_dir: Path, n: int) -> Optional[Dict[str, Any]]: value = _json(state_dir / "iterations" / f"{n:03d}.json", None) - return _public_record(value, _task_spec(state_dir)) if isinstance(value, dict) else None + if not isinstance(value, dict): + return None + return _public_record(state_dir, value, _task_spec(state_dir), _baseline_report(state_dir)) def read_champion(state_dir: Path) -> Dict[str, Any]: - return _json(state_dir / "champion" / "champion.json", {}) or {} + record = _json(state_dir / "champion" / "champion.json", {}) or {} + if not isinstance(record, dict) or not record: + return {} + reference = champion_report_reference(state_dir, record) + benchmark = load_report_reference(state_dir, reference) + promotion_incumbent_ref = record.get("promotion_incumbent_report") + promotion_incumbent = ( + load_report_reference(state_dir, promotion_incumbent_ref) + if isinstance(promotion_incumbent_ref, dict) else {} + ) + profile = _champion_profile(state_dir, record) + return { + **record, + "measurement_report": reference, + "benchmark": benchmark, + "promotion_incumbent": promotion_incumbent, + "profile": profile, + } def read_baseline(state_dir: Path) -> Dict[str, Any]: @@ -43,30 +72,45 @@ def read_baseline(state_dir: Path) -> Dict[str, Any]: initial_hip = _json( state_dir / "certified" / "initial-hip" / "initial-hip-manifest.json", {} ) or {} - profile = _json(state_dir / "system_build" / "build_profile.json", {}) or {} + build_profile = _json(state_dir / "system_build" / "build_profile.json", {}) or {} requirements = _json(state_dir / "requirements.json", {}) or {} - correctness = manifest.get("correctness") or {} - benchmark = manifest.get("benchmark") or {} - hardware_profile = manifest.get("hardware_profile") or {} frozen_profiler = _json( state_dir / "system_profiler" / "profiler_profile.json", {} ) or {} + benchmark = _baseline_report(state_dir) + hardware_profile = _manifest_report( + state_dir, + manifest, + "profile_report", + state_dir / "baseline" / "baseline-hardware-profile.json", + ) + initial_benchmark = _manifest_report( + state_dir, + initial_hip, + "benchmark_report", + state_dir / "certified" / "initial-hip" / "candidate-benchmark-report.json", + ) + initial_profile = _manifest_report( + state_dir, + initial_hip, + "profile_report", + state_dir / "certified" / "initial-hip" / "candidate-hardware-profile.json", + ) spec = _task_spec(state_dir) cases = _baseline_cases(benchmark, spec) - summary = _aggregate(cases, "baseline_ms") return { "certified": bool(manifest), "implementation": manifest.get("implementation", "legacy"), "certified_at": manifest.get("certified_at"), "build_fingerprint": manifest.get("build_fingerprint"), - "backend": profile.get("backend"), - "kernel_language": profile.get("kernel_language"), - "target_hardware": profile.get("target_hardware"), - "gpu_arch": profile.get("gpu_arch"), - "detected_hardware": profile.get("detected_hardware"), - "compiler": profile.get("compiler"), - "compiler_version": profile.get("compiler_version"), - "cmake_version": profile.get("cmake_version"), + "backend": build_profile.get("backend"), + "kernel_language": build_profile.get("kernel_language"), + "target_hardware": build_profile.get("target_hardware"), + "gpu_arch": build_profile.get("gpu_arch"), + "detected_hardware": build_profile.get("detected_hardware"), + "compiler": build_profile.get("compiler"), + "compiler_version": build_profile.get("compiler_version"), + "cmake_version": build_profile.get("cmake_version"), "profiler": { "profile_id": frozen_profiler.get("id"), "tool": frozen_profiler.get("tool_kind"), @@ -76,14 +120,16 @@ def read_baseline(state_dir: Path) -> Dict[str, Any]: "counter_groups": frozen_profiler.get("counter_groups") or [], "passed": hardware_profile.get("passed"), }, - "correctness": correctness.get("summary") or {}, + "correctness": (manifest.get("correctness") or {}).get("summary") or {}, "initial_hip": { "certified": bool(initial_hip), "certified_at": initial_hip.get("certified_at"), "build_fingerprint": initial_hip.get("build_fingerprint"), "correctness": (initial_hip.get("correctness") or {}).get("summary") or {}, - "score": (initial_hip.get("benchmark") or {}).get("score") or {}, - "hardware_profile": initial_hip.get("hardware_profile") or {}, + "score": initial_benchmark.get("score") or {}, + "profile": initial_profile, + "benchmark_report": initial_hip.get("benchmark_report") or {}, + "profile_report": initial_hip.get("profile_report") or {}, }, "task": { "kernel_path": req_field(requirements, "initial_submission"), @@ -93,111 +139,70 @@ def read_baseline(state_dir: Path) -> Dict[str, Any]: }, "benchmark": { "methodology": benchmark.get("methodology") or {}, - "case_count": len(benchmark.get("cases") or []), - "summary": summary, + "case_count": len(cases), + "summary": _measurement_summary(cases), "cases": cases, + "report": manifest.get("benchmark_report") or {}, }, } def read_charts(state_dir: Path) -> Dict[str, Any]: - records = read_iterations(state_dir) - manifest = _json(state_dir / "baseline" / "baseline-manifest.json", {}) or {} spec = _task_spec(state_dir) - baseline_cases = _baseline_cases(manifest.get("benchmark") or {}, spec) - baseline_hardware = manifest.get("hardware_profile") or {} - baseline_summary = _aggregate(baseline_cases, "baseline_ms") + baseline_report = _baseline_report(state_dir) + baseline_manifest = _json( + state_dir / "baseline" / "baseline-manifest.json", {} + ) or {} + baseline_profile = _manifest_report( + state_dir, + baseline_manifest, + "profile_report", + state_dir / "baseline" / "baseline-hardware-profile.json", + ) + baseline_cases = _baseline_cases(baseline_report, spec) + records = read_iterations(state_dir) champion = read_champion(state_dir) - - series: Dict[str, List[Dict[str, Any]]] = { - "latency_ms": [], - "weighted_speedup": [], - "tflops": [], - "bandwidth_gbps": [], - "critical_regression": [], - "duration_s": [], - "measured_bandwidth_gbps": [], - "l2_hit_pct": [], - "compute_busy_pct": [], - "vgpr_count": [], - "lds_bytes": [], - } - if baseline_cases: - _append_point(series["latency_ms"], 0, baseline_summary.get("latency_ms"), True) - _append_point(series["weighted_speedup"], 0, 1.0, True) - _append_point(series["tflops"], 0, baseline_summary.get("tflops"), True) - _append_point( - series["bandwidth_gbps"], 0, baseline_summary.get("bandwidth_gbps"), True - ) - _append_point(series["critical_regression"], 0, 0.0, True) - _append_hardware_points(series, 0, baseline_hardware, True) - - candidate_cases_by_iteration: Dict[int, List[Dict[str, Any]]] = {} + champion_cases = _comparison_cases( + champion.get("promotion_incumbent") or baseline_report, + champion.get("benchmark") or baseline_report, + spec, + ) + champion_profile = champion.get("profile") or baseline_profile + profile_cases = _merge_hardware_cases(champion_cases, champion_profile) + case_series = _case_series( + baseline_cases, + baseline_profile, + records, + spec, + ) + duration_series: List[Dict[str, Any]] = [] for record in records: - iteration = int(record.get("iteration") or 0) - score = record.get("score") or {} - hardware = record.get("hardware_profile") or {} - cases = _score_cases(score.get("cases") or [], spec) - if cases: - candidate_cases_by_iteration[iteration] = cases - summary = _aggregate(cases, "candidate_ms") - _append_point( - series["latency_ms"], iteration, summary.get("latency_ms"), - bool(record.get("promoted")), - ) - _append_point( - series["tflops"], iteration, summary.get("tflops"), - bool(record.get("promoted")), - ) - _append_point( - series["bandwidth_gbps"], iteration, summary.get("bandwidth_gbps"), - bool(record.get("promoted")), - ) - _append_hardware_points( - series, iteration, hardware, bool(record.get("promoted")) - ) _append_point( - series["weighted_speedup"], iteration, score.get("weighted_speedup"), + duration_series, + int(record.get("iteration") or 0), + record.get("duration_s"), bool(record.get("promoted")), ) - _append_point( - series["critical_regression"], iteration, - score.get("critical_regression"), bool(record.get("promoted")), - ) - _append_point( - series["duration_s"], iteration, record.get("duration_s"), - bool(record.get("promoted")), - ) - - champion_iteration = int(champion.get("iteration") or 0) - champion_cases = candidate_cases_by_iteration.get(champion_iteration, baseline_cases) - champion_record = next( - (record for record in records if int(record.get("iteration") or 0) == champion_iteration), - None, - ) - champion_hardware = ( - (champion_record or {}).get("hardware_profile") or baseline_hardware - ) - champion_summary = ( - _aggregate(champion_cases, "candidate_ms") - if champion_iteration in candidate_cases_by_iteration - else baseline_summary - ) - champion_summary = { - **champion_summary, - **_hardware_summary(champion_hardware), - "weighted_speedup": float(champion.get("weighted_speedup", 1.0) or 1.0), - "iteration": champion_iteration, - } return { - "series": series, - "baseline_summary": baseline_summary, - "champion_summary": champion_summary, - "profile_cases": _merge_hardware_cases(champion_cases, champion_hardware), - # Compatibility for early clients of this task-local endpoint. - "weighted_speedup": series["weighted_speedup"], - "critical_regression": series["critical_regression"], - "durations": series["duration_s"], + "series": {"duration_s": duration_series}, + "case_ids": list(case_series), + "case_series": case_series, + "baseline_summary": _measurement_summary(baseline_cases), + "champion_summary": { + **_gate_summary(champion_cases), + "worst_case_speedup": min( + ( + float(case["speedup"]) + for case in champion_cases + if case.get("speedup") is not None + ), + default=None, + ), + "iteration": int(champion.get("iteration") or 0), + "kind": champion.get("kind", "triton"), + "reason": champion.get("reason"), + }, + "profile_cases": profile_cases, } @@ -240,16 +245,98 @@ def _task_spec(state_dir: Path) -> Optional[KernelTaskSpec]: return None -def _public_record(record: Dict[str, Any], spec: Optional[KernelTaskSpec]) -> Dict[str, Any]: +def _manifest_report( + state_dir: Path, + manifest: Dict[str, Any], + key: str, + legacy_path: Path, +) -> Dict[str, Any]: + reference = manifest.get(key) + if isinstance(reference, dict): + return load_report_reference(state_dir, reference) + legacy_key = "benchmark" if key == "benchmark_report" else "hardware_profile" + embedded = manifest.get(legacy_key) + if isinstance(embedded, dict) and embedded: + return embedded + value = _json(legacy_path, {}) + return value if isinstance(value, dict) else {} + + +def _baseline_report(state_dir: Path) -> Dict[str, Any]: + manifest = _json(state_dir / "baseline" / "baseline-manifest.json", {}) or {} + return _manifest_report( + state_dir, + manifest, + "benchmark_report", + state_dir / "baseline" / "baseline-benchmark-report.json", + ) + + +def _champion_profile(state_dir: Path, champion: Dict[str, Any]) -> Dict[str, Any]: + kind = str(champion.get("kind") or "hip") + iteration = int(champion.get("iteration") or 0) + if kind == "triton": + manifest = _json(state_dir / "baseline" / "baseline-manifest.json", {}) or {} + return _manifest_report( + state_dir, + manifest, + "profile_report", + state_dir / "baseline" / "baseline-hardware-profile.json", + ) + if iteration == 0: + manifest = _json( + state_dir / "certified" / "initial-hip" / "initial-hip-manifest.json", {} + ) or {} + return _manifest_report( + state_dir, + manifest, + "profile_report", + state_dir / "certified" / "initial-hip" / "candidate-hardware-profile.json", + ) + record = _json(state_dir / "iterations" / f"{iteration:03d}.json", {}) or {} + reference = record.get("profile_report") + if isinstance(reference, dict): + return load_report_reference(state_dir, reference) + return _json( + state_dir / "logs" / f"{iteration:03d}" / "candidate-hardware-profile.json", + {}, + ) or {} + + +def _public_record( + state_dir: Path, + record: Dict[str, Any], + spec: Optional[KernelTaskSpec], + baseline_report: Dict[str, Any], +) -> Dict[str, Any]: result = copy.deepcopy(record) + measurement_ref = result.get("measurement_report") + if isinstance(measurement_ref, dict): + benchmark = load_report_reference(state_dir, measurement_ref) + result["benchmark"] = benchmark + result["score"] = benchmark.get("score") or {} + else: + benchmark = {"score": result.get("score") or {}} + profile_ref = result.get("profile_report") + if isinstance(profile_ref, dict): + result["profile"] = load_report_reference(state_dir, profile_ref) + elif isinstance(result.get("hardware_profile"), dict): + result["profile"] = result.get("hardware_profile") score = result.get("score") - if not isinstance(score, dict): - return result - score["cases"] = _score_cases(score.get("cases") or [], spec) - private = _private_ids(spec) - score["reasons"] = [ - _redact(str(reason), private) for reason in score.get("reasons") or [] - ] + if isinstance(score, dict): + if not score.get("cases") and benchmark.get("cases"): + score["cases"] = _comparison_cases( + baseline_report, + benchmark, + spec, + ) + else: + score["cases"] = _score_cases(score.get("cases") or [], spec) + private = _private_ids(spec) + score["reasons"] = [ + _redact(str(reason), private) for reason in score.get("reasons") or [] + ] + result.pop("hardware_profile", None) return result @@ -273,12 +360,50 @@ def _baseline_cases( case_id = str(raw.get("id") or "") if not case_id or case_id in private: continue - item = specs.get(case_id) try: latency = float(raw["latency_ms"]) except (KeyError, TypeError, ValueError): continue - cases.append(_profile_case(item, case_id, latency, latency)) + cases.append(_profile_case(specs.get(case_id), case_id, latency, latency)) + return cases + + +def _comparison_cases( + baseline_report: Dict[str, Any], + candidate_report: Dict[str, Any], + spec: Optional[KernelTaskSpec], +) -> List[Dict[str, Any]]: + baseline = { + str(case.get("id")): case + for case in baseline_report.get("cases") or [] + if isinstance(case, dict) and case.get("id") + } + candidate = { + str(case.get("id")): case + for case in candidate_report.get("cases") or [] + if isinstance(case, dict) and case.get("id") + } + specs = _case_specs(spec) + private = _private_ids(spec) + cases: List[Dict[str, Any]] = [] + expected = [case.id for case in spec.benchmark_cases] if spec else list(baseline) + for case_id in expected: + if case_id in private or case_id not in baseline or case_id not in candidate: + continue + try: + baseline_ms = float(baseline[case_id]["latency_ms"]) + candidate_ms = float(candidate[case_id]["latency_ms"]) + except (KeyError, TypeError, ValueError): + continue + shown = _profile_case(specs.get(case_id), case_id, baseline_ms, candidate_ms) + for key in ( + "latency_mean_ms", "latency_median_ms", "latency_stddev_ms", + "latency_cv", "latency_min_ms", "latency_max_ms", + "measurement_batches", "sample_count", + ): + if key in candidate[case_id]: + shown[key] = candidate[case_id][key] + cases.append(shown) return cases @@ -299,23 +424,23 @@ def _score_cases( candidate_ms = float(raw["candidate_ms"]) except (KeyError, TypeError, ValueError): continue - cases.append(_profile_case(specs.get(case_id), case_id, baseline_ms, candidate_ms)) + cases.append(_profile_case( + specs.get(case_id), case_id, baseline_ms, candidate_ms + )) return cases def _profile_case( - spec: Optional[BenchmarkCaseSpec], case_id: str, - baseline_ms: float, candidate_ms: float, + spec: Optional[BenchmarkCaseSpec], + case_id: str, + baseline_ms: float, + candidate_ms: float, ) -> Dict[str, Any]: - weight = float(spec.weight) if spec else 1.0 - critical = bool(spec.critical) if spec else False flops = spec.flops if spec else None transferred = spec.bytes if spec else None return { "id": case_id, "shape": spec.shape if spec else None, - "weight": weight, - "critical": critical, "flops": flops, "bytes": transferred, "baseline_ms": baseline_ms, @@ -329,35 +454,37 @@ def _profile_case( } -def _aggregate(cases: List[Dict[str, Any]], latency_key: str) -> Dict[str, Any]: - valid = [case for case in cases if float(case.get(latency_key) or 0) > 0] - if not valid: - return {"latency_ms": None, "tflops": None, "bandwidth_gbps": None} - total_weight = sum(float(case.get("weight") or 1.0) for case in valid) - weighted_ms = sum( - float(case.get("weight") or 1.0) * float(case[latency_key]) for case in valid - ) - flop_cases = [case for case in valid if case.get("flops") is not None] - byte_cases = [case for case in valid if case.get("bytes") is not None] +def _measurement_summary(cases: List[Dict[str, Any]]) -> Dict[str, Any]: + invalid = [] + for case in cases: + try: + latency_ms = float(case.get("baseline_ms")) + except (TypeError, ValueError): + latency_ms = math.nan + if not math.isfinite(latency_ms) or latency_ms <= 0: + invalid.append(str(case.get("id"))) return { - "latency_ms": weighted_ms / total_weight, - "tflops": _aggregate_rate(flop_cases, latency_key, "flops", 1e9), - "bandwidth_gbps": _aggregate_rate(byte_cases, latency_key, "bytes", 1e6), + "case_count": len(cases), + "all_shapes_measured": bool(cases) and not invalid, + "invalid_case_ids": invalid, } -def _aggregate_rate( - cases: List[Dict[str, Any]], latency_key: str, work_key: str, scale: float, -) -> Optional[float]: - if not cases: - return None - work = sum( - float(case.get("weight") or 1.0) * float(case[work_key]) for case in cases - ) - elapsed = sum( - float(case.get("weight") or 1.0) * float(case[latency_key]) for case in cases - ) - return work / elapsed / scale if elapsed > 0 else None +def _gate_summary(cases: List[Dict[str, Any]]) -> Dict[str, Any]: + return { + "case_count": len(cases), + "all_shapes_passed": bool(cases) and all( + float(case.get("candidate_ms") or 0) + < float(case.get("baseline_ms") or 0) + for case in cases + ), + "failed_case_ids": [ + str(case.get("id")) + for case in cases + if float(case.get("candidate_ms") or 0) + >= float(case.get("baseline_ms") or 0) + ], + } def _rate(work: Optional[float], latency_ms: float, scale: float) -> Optional[float]: @@ -377,33 +504,97 @@ def _append_point( _HARDWARE_METRICS = ( - "measured_bandwidth_gbps", "l2_hit_pct", "compute_busy_pct", - "vgpr_count", "lds_bytes", + "measured_bandwidth_gbps", + "hbm_read_gbps", + "hbm_write_gbps", + "hbm_read_bytes", + "hbm_write_bytes", + "l2_hit_pct", + "occupancy_pct", + "vgpr_count", + "agpr_count", + "sgpr_count", + "lds_bytes", + "scratch_bytes", + "dispatch_count", ) -def _hardware_summary(report: Dict[str, Any]) -> Dict[str, Optional[float]]: - cases = [case for case in report.get("cases") or [] if isinstance(case, dict)] - result: Dict[str, Optional[float]] = {} - for metric in _HARDWARE_METRICS: - values: List[float] = [] +def _case_series( + baseline_cases: List[Dict[str, Any]], + baseline_profile: Dict[str, Any], + records: List[Dict[str, Any]], + spec: Optional[KernelTaskSpec], +) -> Dict[str, Dict[str, Any]]: + baseline_profiled = { + str(case.get("id")): case + for case in _merge_hardware_cases(baseline_cases, baseline_profile) + } + result: Dict[str, Dict[str, Any]] = {} + for case in baseline_cases: + case_id = str(case.get("id") or "") + if not case_id: + continue + series = _empty_case_series() + shown = baseline_profiled.get(case_id, case) + _append_case_points(series, 0, shown, True, True) + result[case_id] = {"case": shown, "series": series} + + for record in records: + iteration = int(record.get("iteration") or 0) + promoted = bool(record.get("promoted")) + cases = _merge_hardware_cases( + _score_cases((record.get("score") or {}).get("cases") or [], spec), + record.get("profile") or {}, + ) for case in cases: - try: - value = float(case[metric]) - except (KeyError, TypeError, ValueError): + case_id = str(case.get("id") or "") + if not case_id or case_id not in result: continue - values.append(value) - result[metric] = sum(values) / len(values) if values else None + _append_case_points( + result[case_id]["series"], iteration, case, promoted, False + ) + _append_point( + result[case_id]["series"]["duration_s"], + iteration, + record.get("duration_s"), + promoted, + ) return result -def _append_hardware_points( - series: Dict[str, List[Dict[str, Any]]], iteration: int, - report: Dict[str, Any], promoted: bool, +def _empty_case_series() -> Dict[str, List[Dict[str, Any]]]: + return { + "latency_ms": [], + "speedup": [], + "tflops": [], + "bandwidth_gbps": [], + "regression": [], + "duration_s": [], + **{metric: [] for metric in _HARDWARE_METRICS}, + } + + +def _append_case_points( + series: Dict[str, List[Dict[str, Any]]], + iteration: int, + case: Dict[str, Any], + promoted: bool, + baseline: bool, ) -> None: - summary = _hardware_summary(report) + prefix = "baseline" if baseline else "candidate" + _append_point(series["latency_ms"], iteration, case.get(f"{prefix}_ms"), promoted) + _append_point(series["speedup"], iteration, 1.0 if baseline else case.get("speedup"), promoted) + _append_point(series["regression"], iteration, 0.0 if baseline else case.get("regression"), promoted) + _append_point(series["tflops"], iteration, case.get(f"{prefix}_tflops"), promoted) + _append_point( + series["bandwidth_gbps"], + iteration, + case.get(f"{prefix}_bandwidth_gbps"), + promoted, + ) for metric in _HARDWARE_METRICS: - _append_point(series[metric], iteration, summary.get(metric), promoted) + _append_point(series[metric], iteration, case.get(metric), promoted) def _merge_hardware_cases( diff --git a/metainfer/tasks/opt_GEMM_kernel/static/gemm-arena-detail.js b/metainfer/tasks/opt_GEMM_kernel/static/gemm-arena-detail.js index c0b6fa5f..b5029f2e 100644 --- a/metainfer/tasks/opt_GEMM_kernel/static/gemm-arena-detail.js +++ b/metainfer/tasks/opt_GEMM_kernel/static/gemm-arena-detail.js @@ -77,7 +77,7 @@ function contractShapes(value) { if (!shapes.length) return "—"; return shapes.map((item) => { const dims = shape(item.shape); - return `${item.id}: ${dims}, weight=${item.weight}${item.critical ? ", critical" : ""}`; + return `${item.id}: ${dims}`; }).join("\n"); } @@ -129,9 +129,24 @@ function GuidancePanel({ taskId, guidance, onSubmitted }) { export default function GemmArenaDetail({ taskId, data }) { const arena = useArena(taskId); + const [selectedCase, setSelectedCase] = useState(""); const nodes = arena.graph?.nodes || []; - const summary = arena.charts?.champion_summary || {}; const cases = arena.charts?.profile_cases || []; + const caseIds = arena.charts?.case_ids || []; + useEffect(() => { + if (!caseIds.length) { + setSelectedCase(""); + } else if (!caseIds.includes(selectedCase)) { + setSelectedCase(caseIds[0]); + } + }, [caseIds.join("\u0000"), selectedCase]); + const selectedProfile = cases.find((item) => item.id === selectedCase) || {}; + const selectedSeries = arena.charts?.case_series?.[selectedCase]?.series || {}; + const championIteration = arena.charts?.champion_summary?.iteration ?? 0; + const chartPayload = { series: selectedSeries }; + const selectedIterationCase = (item) => ( + item.score?.cases || [] + ).find((caseItem) => caseItem.id === selectedCase) || {}; return html`
@@ -152,17 +167,39 @@ export default function GemmArenaDetail({ taskId, data }) {
-

Champion profiler

+
+
+

Champion profiler

+

Every shape must beat a same-round current Champion hipprof trace. Values are raw GPU DurationNs operator means; no shape weighting is used.

+
+ +
-
Speedup${number(summary.weighted_speedup || 1, 4, "×")}
-
Weighted latency${number(summary.latency_ms, 4, " ms")}
-
Compute${number(summary.tflops, 2, " TFLOPS")}
-
Modelled bandwidth${number(summary.bandwidth_gbps, 2, " GB/s")}
-
Measured bandwidth${number(summary.measured_bandwidth_gbps, 2, " GB/s")}
-
L2 hit${number(summary.l2_hit_pct, 2, "%")}
-
Compute busy${number(summary.compute_busy_pct, 2, "%")}
-
Championiter ${summary.iteration ?? 0}
+
Same-round incumbent${number(selectedProfile.baseline_ms, 4, " ms")}
+
Champion latency${number(selectedProfile.candidate_ms, 4, " ms")}
+
Stddev${number(selectedProfile.latency_stddev_ms, 5, " ms")}
+
CV${pct(selectedProfile.latency_cv)}
+
Observed range${number(selectedProfile.latency_min_ms, 4)}–${number(selectedProfile.latency_max_ms, 4, " ms")}
+
Samples / batches${number(selectedProfile.sample_count, 0)} / ${number(selectedProfile.measurement_batches || 1, 0)}
+
Case speedup${number(selectedProfile.speedup, 4, "×")}
+
Regression${pct(selectedProfile.regression)}
+
GPU dispatches/call${number(selectedProfile.dispatch_count, 0)}
+
Compute${number(selectedProfile.candidate_tflops, 2, " TFLOPS")}
+
HBM total${number(selectedProfile.measured_bandwidth_gbps, 2, " GB/s")}
+
HBM read${number(selectedProfile.hbm_read_gbps, 2, " GB/s")}
+
HBM write${number(selectedProfile.hbm_write_gbps, 2, " GB/s")}
+
L2 hit${number(selectedProfile.l2_hit_pct, 2, "%")}
+
Occupancy${number(selectedProfile.occupancy_pct, 2, "%")}
+
Championiter ${championIteration}
+

Promotion: ${arena.charts?.champion_summary?.reason || "certified Triton baseline"}

+

Kernel contribution per operator call: ${compact(selectedProfile.kernel_breakdown_us)}

@@ -183,62 +220,67 @@ export default function GemmArenaDetail({ taskId, data }) {
-

Task contract read-only · frozen evaluator

-
-
Kernel path${arena.baseline.task?.kernel_path || "—"}
-
Data types${compact(arena.baseline.task?.public_contract?.dtype)}
-
Max iterations${arena.baseline.task?.max_iterations ?? "—"}
-
Operation
${arena.baseline.task?.public_contract?.operation || "GEMM"}
-
Public shapes
${contractShapes(arena.baseline.task?.public_contract)}
-
Layout
${compact(arena.baseline.task?.public_contract?.layout)}
-
ABI
${compact(arena.baseline.task?.public_contract?.abi)}
-
-
- -
-

Performance by iteration

- <${ProfilerCharts} payload=${arena.charts} /> +

Performance by iteration · ${selectedCase || "no case"}

+ <${ProfilerCharts} payload=${chartPayload} />

Champion workload profile

- + ${cases.map((item) => html` + + - - - + + - - - - + + `)}
CaseM×N×KLatencySpeedupTFLOPSModelled BWMeasured BWL2 hitCompute busyVGPRLDSCritical
CaseM×N×KSame-round incumbentChampion meanStddev/CVSpeedupDispatchesHBM R/W/totalL2 hitVGPR/AGPR/SGPRLDS/scratch
${item.id} ${shape(item.shape)}${number(item.baseline_ms, 4, " ms")} ${number(item.candidate_ms, 4, " ms")}${number(item.latency_stddev_ms, 5, " ms")} / ${pct(item.latency_cv)} ${number(item.speedup, 3, "×")}${number(item.candidate_tflops, 2)}${number(item.candidate_bandwidth_gbps, 2, " GB/s")}${number(item.measured_bandwidth_gbps, 2, " GB/s")}${number(item.dispatch_count, 0)}${number(item.hbm_read_gbps, 1)} / ${number(item.hbm_write_gbps, 1)} / ${number(item.measured_bandwidth_gbps, 1)} GB/s ${number(item.l2_hit_pct, 2, "%")}${number(item.compute_busy_pct, 2, "%")}${number(item.vgpr_count, 0)}${number(item.lds_bytes, 0, " B")}${item.critical ? "yes" : "no"}${number(item.vgpr_count, 0)} / ${number(item.agpr_count, 0)} / ${number(item.sgpr_count, 0)}${number(item.lds_bytes, 0)} / ${number(item.scratch_bytes, 0)} B
${cases.length ? null : html`

Case metrics appear after baseline certification. TFLOPS requires shape/flops metadata; bandwidth requires bytes metadata in evaluator task.yaml.

`}
-

Iterations

+

Iterations · ${selectedCase || "no case"}

- - ${arena.iterations.map((item) => html` - - - - - - - `)} + + ${arena.iterations.map((item) => { + const caseItem = selectedIterationCase(item); + const failed = item.score?.failed_case_ids || []; + return html` + + + + + + + + + `; + })}
#StatusOutcomeSpeedupCritical regressionPromoted
${item.iteration}${item.status}${item.outcome || "—"}${item.score?.weighted_speedup == null ? "—" : `${Number(item.score.weighted_speedup).toFixed(4)}×`}${pct(item.score?.critical_regression)}${item.promoted ? "yes" : "no"}
#StatusOutcomeCase latencyCase speedupCase regressionAll-shape gatesPromoted
${item.iteration}${item.status}${item.outcome || "—"}${number(caseItem.candidate_ms, 4, " ms")}${number(caseItem.speedup, 4, "×")}${pct(caseItem.regression)}${item.score?.passed ? "passed" : (failed.length ? `failed: ${failed.join(", ")}` : "not run")}${item.promoted ? "yes" : "no"}

Live sub-agents

<${AgentsPanel} agents=${data.agents} />

Event timeline

<${Timeline} events=${data.timeline.events} />
+ +
+

Task contract read-only · frozen evaluator

+
+
Kernel path${arena.baseline.task?.kernel_path || "—"}
+
Data types${compact(arena.baseline.task?.public_contract?.dtype)}
+
Max iterations${arena.baseline.task?.max_iterations ?? "—"}
+
Operation
${arena.baseline.task?.public_contract?.operation || "GEMM"}
+
Public shapes
${contractShapes(arena.baseline.task?.public_contract)}
+
Layout
${compact(arena.baseline.task?.public_contract?.layout)}
+
ABI
${compact(arena.baseline.task?.public_contract?.abi)}
+
+
`; } diff --git a/metainfer/tasks/opt_GEMM_kernel/static/gemm-arena.css b/metainfer/tasks/opt_GEMM_kernel/static/gemm-arena.css index 1f35c2e6..b1c5882b 100644 --- a/metainfer/tasks/opt_GEMM_kernel/static/gemm-arena.css +++ b/metainfer/tasks/opt_GEMM_kernel/static/gemm-arena.css @@ -23,6 +23,10 @@ .gemm-contract-grid code, .gemm-contract-grid pre { display: block; margin: 0; overflow-wrap: anywhere; white-space: pre-wrap; } .gemm-contract-wide { grid-column: 1 / -1; } .gemm-kpis { display: grid; grid-template-columns: repeat(5, minmax(130px, 1fr)); gap: .75rem; } +.gemm-section-heading { align-items: flex-start; display: flex; gap: 1rem; justify-content: space-between; } +.gemm-section-heading h2, .gemm-section-heading p { margin-top: 0; } +.gemm-case-picker { display: grid; gap: .3rem; min-width: min(360px, 40vw); } +.gemm-case-picker span { color: #8b949e; font-size: .75rem; text-transform: uppercase; } .gemm-kpi { border: 1px solid var(--border, #30363d); border-radius: .5rem; padding: .75rem; background: rgba(88, 166, 255, .04); } .gemm-kpi span { color: var(--muted, #8b949e); display: block; font-size: .75rem; margin-bottom: .35rem; text-transform: uppercase; } .gemm-kpi strong { color: #e6edf3; font-size: 1.15rem; } @@ -42,4 +46,6 @@ .gemm-kpis, .gemm-profiler-grid, .gemm-contract-grid { grid-template-columns: 1fr; } .gemm-contract-wide { grid-column: auto; } .gemm-guidance-compose { grid-template-columns: 1fr; } + .gemm-section-heading { display: grid; } + .gemm-case-picker { min-width: 100%; } } diff --git a/metainfer/tasks/opt_GEMM_kernel/static/gemm-profiler-charts.js b/metainfer/tasks/opt_GEMM_kernel/static/gemm-profiler-charts.js index 12e31eb1..ff5e8276 100644 --- a/metainfer/tasks/opt_GEMM_kernel/static/gemm-profiler-charts.js +++ b/metainfer/tasks/opt_GEMM_kernel/static/gemm-profiler-charts.js @@ -5,16 +5,22 @@ import { Chart, registerables } from "chart.js"; Chart.register(...registerables); const DEFINITIONS = [ - ["latency_ms", "Weighted latency", "ms", "#58a6ff"], - ["weighted_speedup", "Speedup vs baseline", "×", "#3fb950"], + ["latency_ms", "Case latency", "ms", "#58a6ff"], + ["speedup", "Case speedup vs baseline", "×", "#3fb950"], ["tflops", "Compute throughput", "TFLOPS", "#d29922"], ["bandwidth_gbps", "Modelled bandwidth", "GB/s", "#f778ba"], - ["measured_bandwidth_gbps", "Profiler memory bandwidth", "GB/s", "#ff9b71"], + ["measured_bandwidth_gbps", "HBM total bandwidth", "GB/s", "#ff9b71"], + ["hbm_read_gbps", "HBM read bandwidth", "GB/s", "#fb7185"], + ["hbm_write_gbps", "HBM write bandwidth", "GB/s", "#fdba74"], ["l2_hit_pct", "L2 hit rate", "%pts", "#2dd4bf"], - ["compute_busy_pct", "Compute busy", "%pts", "#f59e0b"], + ["occupancy_pct", "Reported occupancy", "%pts", "#f59e0b"], ["vgpr_count", "VGPR per work-item", "registers", "#c084fc"], + ["agpr_count", "AGPR per work-item", "registers", "#e879f9"], + ["sgpr_count", "SGPR per wave", "registers", "#a78bfa"], ["lds_bytes", "LDS per workgroup", "bytes", "#22c55e"], - ["critical_regression", "Critical regression", "%", "#a371f7"], + ["scratch_bytes", "Scratch", "bytes", "#84cc16"], + ["dispatch_count", "GPU dispatches per call", "dispatches", "#38bdf8"], + ["regression", "Case regression", "%", "#a371f7"], ["duration_s", "Iteration duration", "s", "#79c0ff"], ]; diff --git a/metainfer/tasks/opt_GEMM_kernel/tests/_helpers.py b/metainfer/tasks/opt_GEMM_kernel/tests/_helpers.py index 0aa1576a..7af7860d 100644 --- a/metainfer/tasks/opt_GEMM_kernel/tests/_helpers.py +++ b/metainfer/tasks/opt_GEMM_kernel/tests/_helpers.py @@ -26,31 +26,37 @@ def make_bundle(root: Path, *, speedup: float = 1.25) -> Path: }, "commands": { phase: {"argv": [sys.executable, "evaluate.py"], "timeout_s": 30} - for phase in ("correctness", "benchmark") + for phase in ("correctness", "profile") }, "cases": { "correctness": ["public", "heldout"], "private": ["heldout"], "benchmark": [ { - "id": "small", "weight": 3, "critical": True, + "id": "small", "shape": {"m": 2, "n": 3, "k": 4, "batch": 1}, "bytes": 100, }, { - "id": "large", "weight": 1, "critical": False, + "id": "large", "shape": {"m": 4, "n": 4, "k": 4, "batch": 2}, "bytes": 200, }, ], }, - "benchmark_protocol": {"warmup": 10, "samples": 100, "timer": "fake"}, - "acceptance": { - "min_weighted_speedup": 1.01, - "noise_threshold": 0.01, - "max_critical_regression": 0.03, - "require_all_cases": True, + "benchmark_protocol": { + "warmup": 10, + "samples": 100, + "trace_calls": 110, + "timer": "hipprof_gpu_kernel_duration_ns", + "statistic": "arithmetic_mean", + "operator_aggregation": "sum_gpu_kernel_duration_per_call", + "synchronization": "hipprof_trace", + "timed_scope": "operator_gpu_dispatches_only", + "host_launch_time_included": False, + "pmc_timing_used": False, }, + "acceptance": {"noise_threshold": 0.01}, } (root / "task.yaml").write_text(yaml.safe_dump(spec), encoding="utf-8") candidate_ms = 1.0 / speedup @@ -142,23 +148,41 @@ def __init__(self): fingerprint="fake-profiler-v1", required=True, ) - def run(self, artifact_dir, output_dir, *, role): + def run( + self, artifact_dir, output_dir, *, role, collection_mode="trace", + case_ids=None, run_label=None, implementation=None, + ): from ..orchestrator.profiler import ProfileResult + latency = ( + 1.0 if role == "baseline" else + (0.95 if "initial-hip" in str(output_dir) else 0.8) + ) report = { "passed": True, "profile_id": "hygon-k100-gfx928", "gpu_arch": "gfx928", - "tool": "rocprofv3", + "tool": "hipprof", "profile_fingerprint": self.profile.fingerprint, "counter_groups": [["SQ_WAVES"]], - "cases": [{ - "id": "small", "vgpr_count": 32, "lds_bytes": 4096, - "l2_hit_pct": 88.0, "compute_busy_pct": 75.0, - "measured_bandwidth_gbps": 640.0, - }], + "cases": [ + { + "id": case_id, "vgpr_count": 32, "lds_bytes": 4096, + "l2_hit_pct": 88.0, "compute_busy_pct": None, + "measured_bandwidth_gbps": 640.0, + } + for case_id in ("small", "large") + ], + "collection_mode": collection_mode, + "timing_cases": [ + {"id": "small", "latency_ms": latency, "kernel_name": "gemm"}, + {"id": "large", "latency_ms": latency, "kernel_name": "gemm"}, + ] if not case_ids else [ + {"id": case_id, "latency_ms": latency, "kernel_name": "gemm"} + for case_id in case_ids + ], } output_dir.mkdir(parents=True, exist_ok=True) - (output_dir / f"{role}-hardware-profile.json").write_text( + (output_dir / f"{run_label or role}-hardware-profile.json").write_text( json.dumps(report), encoding="utf-8" ) return ProfileResult(True, report) diff --git a/metainfer/tasks/opt_GEMM_kernel/tests/test_evaluator.py b/metainfer/tasks/opt_GEMM_kernel/tests/test_evaluator.py index a6307b5c..43a82318 100644 --- a/metainfer/tasks/opt_GEMM_kernel/tests/test_evaluator.py +++ b/metainfer/tasks/opt_GEMM_kernel/tests/test_evaluator.py @@ -17,16 +17,35 @@ def test_runner_executes_all_system_owned_gates(tmp_path): ) assert correctness.passed assert correctness.report["summary"]["expected"] == 2 - baseline = runner.run( - "benchmark", submission, artifacts, tmp_path / "baseline-reports", - role="baseline", build_fingerprint="build-1", + protocol = bundle.spec.benchmark_protocol + baseline = runner.validate_benchmark_report( + { + "passed": True, + "methodology": protocol, + "cases": [ + {"id": "small", "latency_ms": 1.0}, + {"id": "large", "latency_ms": 1.0}, + ], + }, + role="baseline", + build_fingerprint="build-1", ) - benchmark = runner.run( - "benchmark", submission, artifacts, tmp_path / "candidate-reports", - role="candidate", build_fingerprint="build-1", baseline_report=baseline.report, + benchmark = runner.validate_benchmark_report( + { + "passed": True, + "methodology": protocol, + "cases": [ + {"id": "small", "latency_ms": 0.8}, + {"id": "large", "latency_ms": 0.8}, + ], + }, + role="candidate", + build_fingerprint="build-1", + baseline_report=baseline.report, ) + assert baseline.passed assert benchmark.passed - assert benchmark.report["score"]["weighted_speedup"] > 1.2 + assert benchmark.report["score"]["worst_case_speedup"] == 1.25 def test_submission_symlink_is_rejected(tmp_path): diff --git a/metainfer/tasks/opt_GEMM_kernel/tests/test_pipeline.py b/metainfer/tasks/opt_GEMM_kernel/tests/test_pipeline.py index 8d45d955..3bbf0c06 100644 --- a/metainfer/tasks/opt_GEMM_kernel/tests/test_pipeline.py +++ b/metainfer/tasks/opt_GEMM_kernel/tests/test_pipeline.py @@ -5,10 +5,13 @@ from metainfer.orchestrator.state import StateStore from ..orchestrator.evaluator.spec import FrozenEvaluatorBundle -from ..orchestrator.evaluator.champion import ChampionStore +from ..orchestrator.evaluator.champion import ChampionStore, make_report_reference from ..orchestrator.guidance import GuidanceStore from ..orchestrator.phases import graph_payload -from ..orchestrator.pipeline import Orchestrator, OrchestratorConfig +from ..orchestrator.pipeline import ( + Orchestrator, OrchestratorConfig, _combine_hipprof_reports, + _near_promotion_boundary, +) from ._helpers import FakeBuilder, FakeManager, FakeProfiler, make_bundle @@ -77,19 +80,34 @@ def test_one_iteration_promotes_challenger_without_old_task_dependencies(tmp_pat baseline = json.loads( (state / "baseline" / "baseline-manifest.json").read_text(encoding="utf-8") ) - assert baseline["benchmark"]["evaluation_role"] == "baseline" + baseline_report = json.loads( + (state / baseline["benchmark_report"]["path"]).read_text(encoding="utf-8") + ) + baseline_profile = json.loads( + (state / baseline["profile_report"]["path"]).read_text(encoding="utf-8") + ) + assert baseline_report["evaluation_role"] == "baseline" assert baseline["implementation"] == "triton" assert baseline["build_fingerprint"] == "triton-jit" - assert baseline["hardware_profile"]["profile_id"] == "hygon-k100-gfx928" + assert baseline_profile["profile_id"] == "hygon-k100-gfx928" initial_hip = json.loads( (state / "certified" / "initial-hip" / "initial-hip-manifest.json").read_text( encoding="utf-8" ) ) + initial_profile = json.loads( + (state / initial_hip["profile_report"]["path"]).read_text(encoding="utf-8") + ) assert initial_hip["implementation"] == "initial-hip" assert initial_hip["correctness"]["summary"]["expected"] == 2 - assert initial_hip["hardware_profile"]["profile_id"] == "hygon-k100-gfx928" - assert record["hardware_profile"]["cases"][0]["vgpr_count"] == 32 + assert initial_profile["profile_id"] == "hygon-k100-gfx928" + record_profile = json.loads( + (state / record["profile_report"]["path"]).read_text(encoding="utf-8") + ) + assert record_profile["cases"][0]["vgpr_count"] == 32 + assert record["measurement_report"]["path"] == ( + "logs/001/candidate-benchmark-report.json" + ) assert feedback["benchmark"]["hardware_profile"]["gpu_arch"] == "gfx928" assert "Try a 128x128 tile" in manager.prompts["planner"] assert '"entrypoint": "launch_gemm"' in manager.prompts["planner"] @@ -99,8 +117,99 @@ def test_one_iteration_promotes_challenger_without_old_task_dependencies(tmp_pat assert guidance["items"][0]["applied_role"] == "planner" assert any(event["type"] == "human_guidance_applied" for event in timeline) + reloaded = ChampionStore( + state / "champion", + noise_threshold=0.01, + expected_case_ids=["small", "large"], + ).load() + assert reloaded["measurement_report"] == champion["measurement_report"] + champion_report_path = state / champion["measurement_report"]["path"] + original_report = champion_report_path.read_bytes() + champion_report_path.write_text('{"passed": true, "cases": []}', encoding="utf-8") + with pytest.raises(RuntimeError, match="performance report changed"): + ChampionStore( + state / "champion", + noise_threshold=0.01, + expected_case_ids=["small", "large"], + ).load() + champion_report_path.write_bytes(original_report) + (state / "champion" / "submission" / "kernel.cpp").write_text( "// tampered\n", encoding="utf-8" ) with pytest.raises(RuntimeError, match="changed outside promotion"): - ChampionStore(state / "champion", noise_threshold=0.01).load() + ChampionStore( + state / "champion", + noise_threshold=0.01, + expected_case_ids=["small", "large"], + ).load() + + +def test_champion_rejects_candidate_at_exact_noise_boundary(tmp_path): + state = tmp_path / "state" + baseline_path = state / "baseline" / "baseline-benchmark-report.json" + baseline_path.parent.mkdir(parents=True) + baseline_path.write_text( + json.dumps({"cases": [{"id": "shape", "latency_ms": 1.0}]}), + encoding="utf-8", + ) + baseline_ref = make_report_reference(state, baseline_path) + store = ChampionStore( + state / "champion", + noise_threshold=0.01, + expected_case_ids=["shape"], + ) + store.initialize_triton(baseline_ref) + + candidate_dir = tmp_path / "candidate" + candidate_dir.mkdir() + (candidate_dir / "kernel.cpp").write_text("// candidate\n", encoding="utf-8") + candidate_path = state / "logs" / "001" / "candidate-benchmark-report.json" + candidate_path.parent.mkdir(parents=True) + candidate_path.write_text( + json.dumps({"cases": [{"id": "shape", "latency_ms": 0.99}]}), + encoding="utf-8", + ) + candidate_ref = make_report_reference(state, candidate_path) + + promoted, reason, champion = store.consider( + 1, candidate_dir, candidate_ref, baseline_ref, + ) + + assert promoted is False + assert "beyond noise threshold" in reason + assert champion["kind"] == "triton" + assert not store.submission_dir.exists() + + +def test_boundary_retest_combines_raw_hipprof_samples_without_shape_weighting(): + methodology = {"timer": "hipprof_gpu_kernel_duration_ns"} + first = { + "methodology": methodology, + "cases": [{ + "id": "shape", "latency_ms": 1.0, + "operator_samples_ms": [0.9, 1.1], + }], + } + second = { + "methodology": methodology, + "cases": [{ + "id": "shape", "latency_ms": 0.9, + "operator_samples_ms": [0.8, 1.0], + }], + } + combined = _combine_hipprof_reports(first, second) + case = combined["cases"][0] + assert case["latency_ms"] == pytest.approx(0.95) + assert case["sample_count"] == 4 + assert case["measurement_batches"] == 2 + + +def test_one_percent_boundary_triggers_retest(): + incumbent = {"cases": [{"id": "shape", "latency_ms": 1.0}]} + assert _near_promotion_boundary( + incumbent, {"cases": [{"id": "shape", "latency_ms": 0.99}]}, 0.01 + ) + assert not _near_promotion_boundary( + incumbent, {"cases": [{"id": "shape", "latency_ms": 0.8}]}, 0.01 + ) diff --git a/metainfer/tasks/opt_GEMM_kernel/tests/test_plugin.py b/metainfer/tasks/opt_GEMM_kernel/tests/test_plugin.py index 2fc8f5b1..43c1841d 100644 --- a/metainfer/tasks/opt_GEMM_kernel/tests/test_plugin.py +++ b/metainfer/tasks/opt_GEMM_kernel/tests/test_plugin.py @@ -61,53 +61,128 @@ def test_task_does_not_import_other_task_packages(): def test_profiler_chart_payload_uses_frozen_work_metadata(tmp_path): + from ..orchestrator.evaluator.champion import ( + make_report_reference, + write_json_atomic, + ) + state = tmp_path / "state" evaluator = state / "system_evaluator" shutil.copytree(make_bundle(tmp_path / "bundle"), evaluator) (state / "baseline").mkdir(parents=True) (state / "iterations").mkdir(parents=True) (state / "champion").mkdir(parents=True) - baseline = { - "benchmark": { - "methodology": {"warmup": 10, "samples": 100, "timer": "fake"}, + (state / "logs" / "001").mkdir(parents=True) + + protocol = { + "warmup": 10, + "samples": 100, + "trace_calls": 110, + "timer": "hipprof_gpu_kernel_duration_ns", + "statistic": "arithmetic_mean", + "operator_aggregation": "sum_gpu_kernel_duration_per_call", + "synchronization": "hipprof_trace", + "timed_scope": "operator_gpu_dispatches_only", + "host_launch_time_included": False, + "pmc_timing_used": False, + } + baseline_report_path = state / "baseline" / "baseline-benchmark-report.json" + baseline_profile_path = state / "baseline" / "baseline-hardware-profile.json" + candidate_report_path = state / "logs" / "001" / "candidate-benchmark-report.json" + candidate_profile_path = state / "logs" / "001" / "candidate-hardware-profile.json" + write_json_atomic(baseline_report_path, { + "passed": True, + "methodology": protocol, + "cases": [ + {"id": "small", "latency_ms": 2.0}, + {"id": "large", "latency_ms": 4.0}, + ], + }) + write_json_atomic(baseline_profile_path, { + "passed": True, + "cases": [], + }) + write_json_atomic(candidate_report_path, { + "passed": True, + "methodology": protocol, + "cases": [ + {"id": "small", "latency_ms": 1.0}, + {"id": "large", "latency_ms": 2.0}, + ], + "score": { + "passed": True, + "worst_case_speedup": 2.0, + "failed_case_ids": [], "cases": [ - {"id": "small", "latency_ms": 2.0}, - {"id": "large", "latency_ms": 4.0}, + {"id": "small", "baseline_ms": 2.0, "candidate_ms": 1.0}, + {"id": "large", "baseline_ms": 4.0, "candidate_ms": 2.0}, ], }, - } + }) + write_json_atomic(candidate_profile_path, { + "passed": True, + "cases": [{ + "id": "small", "vgpr_count": 40, "agpr_count": 8, + "sgpr_count": 32, "lds_bytes": 8192, "scratch_bytes": 0, + "l2_hit_pct": 91.0, "occupancy_pct": None, + "measured_bandwidth_gbps": 700.0, + "hbm_read_gbps": 600.0, "hbm_write_gbps": 100.0, + "dispatch_count": 2, + "kernel_breakdown_us": {"splitk": 0.8, "reduce": 0.2}, + }], + }) + baseline_ref = make_report_reference(state, baseline_report_path) + baseline_profile_ref = make_report_reference(state, baseline_profile_path) + candidate_ref = make_report_reference(state, candidate_report_path) + candidate_profile_ref = make_report_reference(state, candidate_profile_path) (state / "baseline" / "baseline-manifest.json").write_text( - json.dumps(baseline), encoding="utf-8" + json.dumps({ + "implementation": "triton", + "benchmark_report": baseline_ref, + "profile_report": baseline_profile_ref, + }), + encoding="utf-8", ) record = { "iteration": 1, "duration_s": 12, "promoted": True, - "score": { - "weighted_speedup": 2.0, - "critical_regression": 0.0, - "cases": [ - {"id": "small", "baseline_ms": 2.0, "candidate_ms": 1.0, - "flops": 1, "bytes": 1}, - {"id": "large", "baseline_ms": 4.0, "candidate_ms": 2.0}, - ], - }, - "hardware_profile": { - "cases": [{"id": "small", "vgpr_count": 40, "lds_bytes": 8192, - "l2_hit_pct": 91.0, "compute_busy_pct": 82.0, - "measured_bandwidth_gbps": 700.0}], - }, + "measurement_report": candidate_ref, + "profile_report": candidate_profile_ref, } (state / "iterations" / "001.json").write_text(json.dumps(record), encoding="utf-8") (state / "champion" / "champion.json").write_text( - json.dumps({"iteration": 1, "weighted_speedup": 2.0}), encoding="utf-8" + json.dumps({ + "schema_version": 2, + "kind": "hip", + "iteration": 1, + "measurement_report": candidate_ref, + "reason": "every shape passed baseline and Champion noise gates", + }), + encoding="utf-8", ) payload = _state_readers.read_charts(state) - assert payload["champion_summary"]["weighted_speedup"] == 2.0 + assert payload["baseline_summary"] == { + "case_count": 2, + "all_shapes_measured": True, + "invalid_case_ids": [], + } + assert payload["champion_summary"]["all_shapes_passed"] is True + assert payload["champion_summary"]["worst_case_speedup"] == 2.0 + assert "weighted_speedup" not in payload + assert "critical_regression" not in payload assert payload["profile_cases"][0]["candidate_tflops"] == 48 / 1.0 / 1e9 assert payload["profile_cases"][0]["candidate_bandwidth_gbps"] == 100 / 1.0 / 1e6 - assert payload["series"]["latency_ms"][0]["x"] == 0 - assert payload["series"]["latency_ms"][1]["x"] == 1 - assert payload["series"]["measured_bandwidth_gbps"][0]["y"] == 700.0 + assert payload["profile_cases"][0]["measured_bandwidth_gbps"] == 700.0 assert payload["profile_cases"][0]["vgpr_count"] == 40 + assert payload["profile_cases"][0]["dispatch_count"] == 2 + assert payload["profile_cases"][0]["kernel_breakdown_us"]["reduce"] == 0.2 + assert payload["case_ids"] == ["small", "large"] + small = payload["case_series"]["small"]["series"] + assert [point["y"] for point in small["latency_ms"]] == [2.0, 1.0] + assert [point["y"] for point in small["speedup"]] == [1.0, 2.0] + assert small["tflops"][1]["y"] == 48 / 1.0 / 1e9 + assert small["measured_bandwidth_gbps"][0]["x"] == 1 + large = payload["case_series"]["large"]["series"] + assert [point["y"] for point in large["latency_ms"]] == [4.0, 2.0] diff --git a/metainfer/tasks/opt_GEMM_kernel/tests/test_profiler.py b/metainfer/tasks/opt_GEMM_kernel/tests/test_profiler.py index 4cde8c06..721b586f 100644 --- a/metainfer/tasks/opt_GEMM_kernel/tests/test_profiler.py +++ b/metainfer/tasks/opt_GEMM_kernel/tests/test_profiler.py @@ -1,6 +1,7 @@ from __future__ import annotations import csv +import importlib.util import subprocess from pathlib import Path @@ -9,7 +10,7 @@ from ..orchestrator.hardware import HardwareProfileError, require_hardware_profile from ..orchestrator.profiler import ( FrozenProfilerProfile, ProfilerError, ProfilerRunner, _parse_case, - _validate_harness_profile, _version, + _python_executable, _validate_harness_profile, _version, ) @@ -48,6 +49,14 @@ def fake_run(argv, **_kwargs): ) +def test_python_executable_ignores_launcher_argv0(monkeypatch): + monkeypatch.setattr( + "metainfer.tasks.opt_GEMM_kernel.orchestrator.profiler.sys.executable", + "/usr/local/bin/metainfer-orchestrator", + ) + assert Path(_python_executable()).name.startswith("python") + + def test_rocprof_csv_is_normalized_for_ui_and_f_agent(tmp_path: Path): path = tmp_path / "pass_1" / "results.csv" path.parent.mkdir() @@ -141,3 +150,92 @@ def test_harness_profile_must_match_requested_case(tmp_path: Path): _validate_harness_profile(tmp_path, "wq-b-tp4-m16") with pytest.raises(ProfilerError, match="does not match"): _validate_harness_profile(tmp_path, "wq-b-tp4-m1") + + +def _analyzer_module(): + path = ( + Path(__file__).resolve().parents[1] + / "harness" / "user_gemm" / "analyze_hipprof_suite.py" + ) + spec = importlib.util.spec_from_file_location("gemm_hipprof_analyzer", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_trace_operator_time_sums_all_dispatches_per_final_sample(): + analyzer = _analyzer_module() + case = { + "id": "split-k", + "host_epoch_begin_ns": 100, + "host_epoch_end_ns": 1000, + } + rows = [ + {"begin_ns": 50, "duration_ns": 9999, "kernel_name": "prepare"}, + {"begin_ns": 100, "duration_ns": 900, "kernel_name": "split"}, + {"begin_ns": 110, "duration_ns": 100, "kernel_name": "split"}, + {"begin_ns": 120, "duration_ns": 100, "kernel_name": "reduce"}, + {"begin_ns": 200, "duration_ns": 100, "kernel_name": "split"}, + {"begin_ns": 210, "duration_ns": 200, "kernel_name": "split"}, + {"begin_ns": 220, "duration_ns": 50, "kernel_name": "reduce"}, + {"begin_ns": 300, "duration_ns": 150, "kernel_name": "split"}, + {"begin_ns": 310, "duration_ns": 250, "kernel_name": "split"}, + {"begin_ns": 320, "duration_ns": 100, "kernel_name": "reduce"}, + {"begin_ns": 1100, "duration_ns": 9999, "kernel_name": "post"}, + ] + operator_us, breakdown, dispatches, samples = analyzer._trace_case_times( + rows, case, calls=3, samples=2 + ) + assert operator_us == pytest.approx(0.425) + assert breakdown == pytest.approx({"split": 0.35, "reduce": 0.075}) + assert dispatches == 3 + assert samples == pytest.approx([0.35, 0.5]) + + +def test_trace_rejects_unstable_final_dispatch_pattern(): + analyzer = _analyzer_module() + case = { + "id": "unstable", + "host_epoch_begin_ns": 100, + "host_epoch_end_ns": 1000, + } + rows = [ + {"begin_ns": 100, "duration_ns": 10, "kernel_name": "main"}, + {"begin_ns": 110, "duration_ns": 10, "kernel_name": "reduce"}, + {"begin_ns": 200, "duration_ns": 10, "kernel_name": "main"}, + {"begin_ns": 210, "duration_ns": 10, "kernel_name": "other"}, + ] + with pytest.raises(RuntimeError, match="unstable measured dispatch pattern"): + analyzer._trace_case_times(rows, case, calls=2, samples=2) + + +def test_pmc_is_normalized_per_call_and_replay_duration_is_not_latency(): + analyzer = _analyzer_module() + case = { + "id": "pmc", + "host_monotonic_begin_ns": 100, + "host_monotonic_end_ns": 1000, + } + rows = [ + { + "BeginNs": str(begin), + "DurationNs": str(duration), + "KernelName": kernel, + "TCC_EA_RDREQ[0]": "1", + "TCC_EA_RDREQ_32B[0]": "1", + "TCC_EA_WRREQ[0]": "1", + "TCC_EA_WRREQ_64B[0]": "0", + } + for begin, duration, kernel in ( + (100, 1_000_000, "split"), + (110, 2_000_000, "reduce"), + (200, 3_000_000, "split"), + (210, 4_000_000, "reduce"), + ) + ] + counters = analyzer._aggregate_counters(rows, case, calls=2) + assert counters["dispatch_count"] == 2 + assert counters["hbm_read_bytes"] == 64 + assert counters["hbm_write_bytes"] == 64 + assert "duration" not in counters diff --git a/metainfer/tasks/opt_GEMM_kernel/tests/test_scoring.py b/metainfer/tasks/opt_GEMM_kernel/tests/test_scoring.py index 704904a3..1efa9bb3 100644 --- a/metainfer/tasks/opt_GEMM_kernel/tests/test_scoring.py +++ b/metainfer/tasks/opt_GEMM_kernel/tests/test_scoring.py @@ -1,50 +1,58 @@ -from ..orchestrator.evaluator.scoring import compare_measurements, score_benchmark +from ..orchestrator.evaluator.scoring import ( + compare_against_champion, + compare_measurements, +) from ..orchestrator.evaluator.spec import AcceptanceSpec, BenchmarkCaseSpec -def test_trace_weighted_score_and_critical_gate(): - result = score_benchmark( - [ - {"id": "hot", "baseline_ms": 10, "candidate_ms": 5, "weight": 9, "critical": True}, - {"id": "cold", "baseline_ms": 10, "candidate_ms": 20, "weight": 1}, - ], - ["hot", "cold"], - AcceptanceSpec(min_weighted_speedup=1.2, max_critical_regression=0.03), +def _spec(case_id: str) -> BenchmarkCaseSpec: + return BenchmarkCaseSpec( + case_id, + shape={"m": 1, "n": 1, "k": 1, "batch": 1}, ) - assert result.passed - assert result.weighted_speedup == 100 / 65 -def test_missing_shape_is_a_hard_failure(): - result = score_benchmark( - [{"id": "a", "baseline_ms": 1, "candidate_ms": 0.5}], - ["a", "b"], +def test_every_shape_must_beat_the_frozen_baseline(): + result = compare_measurements( + [ + {"id": "hot", "latency_ms": 10.0}, + {"id": "cold", "latency_ms": 10.0}, + ], + [ + {"id": "hot", "latency_ms": 5.0}, + {"id": "cold", "latency_ms": 20.0}, + ], + [_spec("hot"), _spec("cold")], AcceptanceSpec(), ) assert not result.passed - assert result.missing_case_ids == ["b"] + assert result.failed_case_ids == ["cold"] + assert result.worst_case_speedup == 0.5 -def test_critical_regression_blocks_good_average(): - result = score_benchmark( +def test_missing_shape_is_a_hard_failure(): + result = compare_measurements( [ - {"id": "hot", "baseline_ms": 100, "candidate_ms": 50, "weight": 10}, - {"id": "critical", "baseline_ms": 1, "candidate_ms": 1.1, "weight": 1, "critical": True}, + {"id": "a", "latency_ms": 1.0}, + {"id": "b", "latency_ms": 1.0}, ], - ["hot", "critical"], - AcceptanceSpec(max_critical_regression=0.03), + [{"id": "a", "latency_ms": 0.5}], + [_spec("a"), _spec("b")], + AcceptanceSpec(), ) assert not result.passed - assert result.critical_regression > 0.09 + assert result.missing_case_ids == ["b"] def test_non_finite_latency_is_rejected(): - result = score_benchmark( - [{"id": "a", "baseline_ms": 1, "candidate_ms": float("nan")}], - ["a"], + result = compare_measurements( + [{"id": "a", "latency_ms": 1.0}], + [{"id": "a", "latency_ms": float("nan")}], + [_spec("a")], AcceptanceSpec(), ) assert not result.passed + assert any("positive" in reason for reason in result.reasons) def test_profiler_rates_are_derived_from_frozen_case_spec(): @@ -52,7 +60,7 @@ def test_profiler_rates_are_derived_from_frozen_case_spec(): [{"id": "gemm", "latency_ms": 2.0}], [{"id": "gemm", "latency_ms": 1.0, "flops": 1}], [BenchmarkCaseSpec( - "gemm", weight=1.0, critical=True, + "gemm", shape={"m": 1000, "n": 1000, "k": 1000, "batch": 1}, flops=2_000_000_000.0, bytes=1_000_000_000.0, @@ -62,3 +70,32 @@ def test_profiler_rates_are_derived_from_frozen_case_spec(): case = result.cases[0] assert case["candidate_tflops"] == 2.0 assert case["candidate_bandwidth_gbps"] == 1000.0 + + +def test_champion_noise_gate_requires_every_shape_to_cross_threshold(): + result = compare_against_champion( + [ + {"id": "a", "latency_ms": 1.0}, + {"id": "b", "latency_ms": 2.0}, + ], + [ + {"id": "a", "latency_ms": 0.98}, + {"id": "b", "latency_ms": 1.99}, + ], + ["a", "b"], + 0.01, + ) + assert not result.passed + assert result.failed_case_ids == ["b"] + + +def test_strict_baseline_gate_rejects_equal_latency(): + result = compare_against_champion( + [{"id": "a", "latency_ms": 1.0}], + [{"id": "a", "latency_ms": 1.0}], + ["a"], + 0.0, + strict=True, + ) + assert not result.passed + assert result.failed_case_ids == ["a"] diff --git a/metainfer/tasks/opt_GEMM_kernel/tests/test_spec.py b/metainfer/tasks/opt_GEMM_kernel/tests/test_spec.py index ba9bd4ad..274acbe1 100644 --- a/metainfer/tasks/opt_GEMM_kernel/tests/test_spec.py +++ b/metainfer/tasks/opt_GEMM_kernel/tests/test_spec.py @@ -12,16 +12,63 @@ def test_task_owned_harness_starter_has_a_valid_protocol(): harness = Path(__file__).resolve().parents[1] / "harness" / "user_gemm" spec = KernelTaskSpec.load(harness / "task.yaml") assert spec.name == "deepseek-w8a8-gemm-tp4-tp8" - assert set(spec.commands) == {"correctness", "benchmark", "profile"} + assert set(spec.commands) == {"correctness", "profile"} assert len(spec.benchmark_cases) == 60 assert len(spec.correctness_case_ids) == 64 assert len(spec.private_case_ids) == 4 assert spec.agent_contract()["abi"]["entrypoint"] == "launch_w8a8_gemm" + assert spec.benchmark_protocol["timer"] == "hipprof_gpu_kernel_duration_ns" + assert spec.benchmark_protocol["operator_aggregation"] == ( + "sum_gpu_kernel_duration_per_call" + ) + assert spec.benchmark_protocol["pmc_timing_used"] is False assert {case.shape["m"] for case in spec.benchmark_cases} == {1, 2, 4, 8, 16, 4096} assert {case.shape["batch"] for case in spec.benchmark_cases} == {1} - small = sum(case.weight for case in spec.benchmark_cases if case.shape["m"] <= 16) - large = sum(case.weight for case in spec.benchmark_cases if case.shape["m"] == 4096) - assert small == pytest.approx(large) + assert all( + set(item) == {"id", "shape"} + for item in spec.agent_contract()["benchmark_shapes"] + ) + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("timer", "gpu_event"), + ("statistic", "median"), + ("operator_aggregation", "longest_kernel"), + ("synchronization", "gpu_event"), + ("timed_scope", "host_api_call"), + ("host_launch_time_included", True), + ("pmc_timing_used", True), + ], +) +def test_benchmark_protocol_rejects_non_hipprof_timing(tmp_path, field, value): + import yaml + + source = make_bundle(tmp_path / "source") + task_path = source / "task.yaml" + raw = yaml.safe_load(task_path.read_text(encoding="utf-8")) + raw["benchmark_protocol"][field] = value + task_path.write_text(yaml.safe_dump(raw), encoding="utf-8") + + with pytest.raises(SpecError, match=rf"benchmark_protocol\.{field}"): + KernelTaskSpec.load(task_path) + + +def test_task_owned_hipprof_suite_is_self_contained(): + harness = Path(__file__).resolve().parents[1] / "harness" / "user_gemm" + collector = (harness / "run_hipprof_suite.py").read_text(encoding="utf-8") + analyzer = (harness / "analyze_hipprof_suite.py").read_text(encoding="utf-8") + evaluator = (harness / "evaluate.py").read_text(encoding="utf-8") + assert "/data/FF/kernel benchmark" not in collector + analyzer + assert "METAINFER_BUILD_ARTIFACT_DIR" in collector + assert "METAINFER_WEIGHT_BUNDLE" in collector + assert '"--hip-trace", "--stats"' in collector + assert '"--pmc-read"' in collector and '"--pmc-write"' in collector + assert 'phase == "profile-batch"' in evaluator + assert '"host_epoch_begin_ns"' in evaluator + assert '"host_epoch_end_ns"' in evaluator + assert '"host_epoch_begin_ns" in case' in analyzer def test_mygemm_baseline_is_decoupled_from_harness_code(): diff --git a/metainfer/tasks/sglang_trace_analyze/__init__.py b/metainfer/tasks/sglang_trace_analyze/__init__.py new file mode 100644 index 00000000..afb0e3af --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/__init__.py @@ -0,0 +1,8 @@ +"""sglang_trace_analyze — auto-generate torch profiler traces via SGLang, +analyze them (operator-to-structure mapping, kernel hotspots, TFLOPS / MFU, +overlap opportunities, fuse suggestions), and surface results + LLM hints +in the MetaInfer WebUI. +""" + +from .orchestrator import plugin as _task_plugin # noqa: F401 +from .server import plugin as _web_plugin # noqa: F401 diff --git a/metainfer/tasks/sglang_trace_analyze/form.yaml b/metainfer/tasks/sglang_trace_analyze/form.yaml new file mode 100644 index 00000000..fdd2dc1b --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/form.yaml @@ -0,0 +1,67 @@ +- key: model_path + header: Model Path + question: "HuggingFace repo id or local path to the model weights." + required: true + form: text + +- key: version + header: Version + question: "Short tag for this run — used in trace directory naming and result labels." + required: true + form: text + +- key: batch_sizes + header: Batch Sizes + question: "Comma-separated list of decode batch sizes to profile, e.g. 1,4,8,16." + required: true + form: text + +- key: mapping_batch_size + header: Mapping BS + question: "Batch size for the mapping run (CUDA Graph disabled). One value is enough — kernel-to-layer mapping is independent of batch size." + required: true + default: "8" + form: number + +- key: input_len + header: Input Len + question: "Synthetic input sequence length." + required: true + default: "512" + form: number + +- key: output_len + header: Output Len + question: "Synthetic output sequence length." + required: true + default: "2000" + form: number + +- key: tp_size + header: TP Size + question: "Tensor-parallelism degree." + required: true + default: "1" + form: number + +- key: pp_size + header: PP Size + question: "Pipeline-parallelism degree." + required: true + default: "1" + form: number + +- key: gpu_model + header: GPU Model + question: "GPU model — used to look up theoretical peak TFLOPS and memory bandwidth." + required: true + form: select + options: + - label: "K100" + description: "FP32 49TF, TF32 98TF, BF16/FP16 192TF, INT8 392TOPS, BW 700GB/s" + - label: "A100_80G" + description: "FP32 19.5TF, TF32 156TF, BF16/FP16 312TF, INT8 624TOPS, BW 2039GB/s" + - label: "H100" + description: "FP32 67TF, TF32 989TF, BF16/FP16 989TF, INT8 1979TOPS, BW 3350GB/s" + - label: "B200" + description: "FP32 90TF, TF32 2250TF, BF16/FP16 2250TF, INT8 4500TOPS, BW 8000GB/s" diff --git a/metainfer/tasks/sglang_trace_analyze/orchestrator/__init__.py b/metainfer/tasks/sglang_trace_analyze/orchestrator/__init__.py new file mode 100644 index 00000000..43d16d09 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/orchestrator/__init__.py @@ -0,0 +1,6 @@ +"""Orchestrator (worker subprocess) for sglang_trace_analyze.""" + +from metainfer.orchestrator.tasks import register +from .plugin import PLUGIN + +register(PLUGIN) diff --git a/metainfer/tasks/sglang_trace_analyze/orchestrator/cli.py b/metainfer/tasks/sglang_trace_analyze/orchestrator/cli.py new file mode 100644 index 00000000..d4179885 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/orchestrator/cli.py @@ -0,0 +1,46 @@ +"""CLI entry point for the sglang_trace_analyze orchestrator subprocess. + +The launcher spawns:: + + python -m run --state-dir … --workspace-dir … + +Contract required by the framework (§6d): ``run`` subcommand + ``--state-dir`` +and ``--workspace-dir`` flags. +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(prog="metainfer-orchestrator") + sub = parser.add_subparsers(dest="command") + + run_p = sub.add_parser("run") + run_p.add_argument("requirements", type=Path, + help="Path to requirements.json") + run_p.add_argument("--state-dir", type=Path, required=True) + run_p.add_argument("--workspace-dir", type=Path, required=True) + # Task-specific flags + run_p.add_argument("--iter-limit", type=int, default=None, + help="Override max iterations (default: derive from batch count)") + + args = parser.parse_args(argv) + if args.command != "run": + parser.print_help() + return 1 + + from .orchestrator import run_with_requirements + return run_with_requirements( + requirements_path=args.requirements, + state_dir=args.state_dir, + workspace_dir=args.workspace_dir, + iter_limit=args.iter_limit, + ) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/metainfer/tasks/sglang_trace_analyze/orchestrator/flops_calculator.py b/metainfer/tasks/sglang_trace_analyze/orchestrator/flops_calculator.py new file mode 100644 index 00000000..288d4a23 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/orchestrator/flops_calculator.py @@ -0,0 +1,195 @@ +"""Compute TFLOPS, bandwidth, and MFU for aggregated kernel entries. + +Uses: +- ``gpu_specs.py`` for theoretical peak values +- kernel ``input_dims`` (from MAPPING trace) or shape rules (for CUDA Graph + formal traces) to derive actual FLOP counts per invocation +- kernel ``total_dur_us`` to compute actual TFLOPS/bandwidth +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +from .gpu_specs import GpuSpec + + +def extract_ck_tile_dims(kernel_name: str) -> tuple | None: + """Extract (M, N, K) tile dimensions from a CK GEMM kernel name. + + Example: ``Cijk_Alik_Bljk_SB_MT64x128x16_...`` → (64, 128, 16) + """ + import re + m = re.search(r"MT(\d+)x(\d+)x(\d+)", kernel_name) + if m: + return int(m.group(1)), int(m.group(2)), int(m.group(3)) + return None + + +def calculate_mfu( + kernels: List[Dict[str, Any]], + gpu_spec: GpuSpec, + *, + batch_size: int, + dtype: str = "bf16", +) -> List[Dict[str, Any]]: + """Augment each kernel entry with TFLOPS, bandwidth, MFU, and bound classification. + + Args: + kernels: Aggregated kernel list. Each entry must have ``total_dur_us`` + and ``count``. Entries from a non-CUDA Graph trace may also have + ``input_dims``, which are used for FLOP/byte estimation where available. + gpu_spec: GPU theoretical peak specification. + batch_size: Decode batch size used for this trace. + dtype: Compute dtype — determines which TFLOPS peak to use. + One of ``fp32``, ``tf32``, ``bf16``, ``fp16``, ``int8``. + + Returns: + The same kernel list with added fields: ``tflops_actual``, + ``bandwidth_gb_s``, ``mfu``, ``bound``, ``flops_per_invocation``. + """ + theoretical_tflops = _theoretical_peak(gpu_spec, dtype) + theoretical_bw = gpu_spec.bandwidth_gb_s + + for k in kernels: + dur_s = k["total_dur_us"] / 1e6 + count = k.get("count", 1) + dur_per_invocation_s = dur_s / count if count else dur_s + dims = k.get("input_dims", []) + op_type = k.get("op_type", "Other") + kernel_name = k.get("kernel_name", "") + + flops = _estimate_flops(op_type, dims, batch_size) + bytes_moved = _estimate_bytes(op_type, dims, batch_size) + + # For CK GEMM kernels without input dims, estimate from tile name + if flops == 0 and op_type == "GEMM": + tile = extract_ck_tile_dims(kernel_name) + if tile: + M, N, K_tile = tile + flops = 2 * M * N * K_tile * count + bytes_moved = (M * K_tile + K_tile * N + M * N) * 2 * count + + tflops_actual = (flops / dur_s / 1e12) if dur_s > 0 else 0 + bandwidth_gb_s = (bytes_moved / dur_s / 1e9) if dur_s > 0 else 0 + mfu = (tflops_actual / theoretical_tflops * 100) if theoretical_tflops > 0 else 0 + + # Compute-bound vs memory-bound heuristic + ops_per_byte = flops / bytes_moved if bytes_moved > 0 else float("inf") + if theoretical_bw > 0: + crossover = theoretical_tflops * 1e12 / (theoretical_bw * 1e9) + else: + crossover = float("inf") + bound = "compute" if ops_per_byte > crossover else "memory" + + k["tflops_actual"] = round(tflops_actual, 6) if flops > 0 else None + k["tflops_theoretical"] = theoretical_tflops + k["bandwidth_gb_s"] = round(bandwidth_gb_s, 1) if bytes_moved > 0 else None + k["bandwidth_theoretical"] = theoretical_bw + k["mfu"] = round(mfu, 3) if flops > 0 else None + k["bound"] = bound if (flops > 0 and bytes_moved > 0) else "unknown" + k["flops_per_invocation"] = int(flops) + + return kernels + + +def _theoretical_peak(spec: GpuSpec, dtype: str) -> float: + """Return theoretical peak TFLOPS for the given dtype.""" + return { + "fp32": spec.fp32_tflops, + "tf32": spec.tf32_tflops, + "bf16": spec.bf16_tflops, + "fp16": spec.fp16_tflops, + "int8": spec.int8_tops, # TOPS → TFLOPS approximate + }.get(dtype, spec.bf16_tflops) + + +def _estimate_flops( + op_type: str, + dims: List[Any], + batch_size: int, +) -> float: + """Estimate FLOPs for one kernel invocation. + + For GEMM: 2*M*N*K (or 2*B*M*N*K for batched). + For Attention: approximately 4*B*seq_len*head_dim*num_heads^2. + For ElementWise: 2*num_elements. + + Returns 0 if dims are unavailable (CUDA Graph trace). + """ + if not dims: + return 0 + + # Use the first observed dim list + d = dims[0] + + if op_type == "GEMM": + if isinstance(d, list) and len(d) >= 2: + if len(d) == 3: + M, K, N = int(d[0]), int(d[1]), int(d[2]) + return 2 * M * N * K + B, M, N, K = _unpack_4d(d, batch_size) + return 2 * B * M * N * K + + elif op_type == "Attention": + if isinstance(d, list) and len(d) >= 3: + seq_len = int(d[0]) + num_heads = int(d[1]) + head_dim = int(d[2]) + return 4 * seq_len * head_dim * num_heads * num_heads * batch_size + + elif op_type == "MoE": + if isinstance(d, list) and len(d) >= 3: + M, K, N = int(d[0]), int(d[1]), int(d[2]) + return 2 * M * N * K + + return 0 + + +def _estimate_bytes( + op_type: str, + dims: List[Any], + batch_size: int, +) -> float: + """Estimate bytes moved (reads + writes) for one kernel invocation. + + Simple heuristic: for GEMM, input_bytes ≈ (M*K + K*N) * dtype_size, + output_bytes ≈ M*N * dtype_size. For elementwise, ≈ 3 * num_elements. + + Returns 0 if dims are unavailable. + """ + if not dims: + return 0 + + d = dims[0] + dtype_size = 2 # bf16/fp16 default + + if op_type == "GEMM": + if isinstance(d, list): + if len(d) == 3: + M, K, N = int(d[0]), int(d[1]), int(d[2]) + return (M * K + K * N + M * N) * dtype_size + B, M, N, K = _unpack_4d(d, batch_size) + return B * (M * K + K * N + M * N) * dtype_size + + elif op_type == "Attention": + if isinstance(d, list) and len(d) >= 3: + seq_len = int(d[0]) + num_heads = int(d[1]) + head_dim = int(d[2]) + # Q, K, V reads + output write (approximate) + return batch_size * seq_len * num_heads * head_dim * 4 * dtype_size + + return 0 + + +def _unpack_4d( + dims: list, + batch_size: int, +) -> tuple: + """Unpack a 4-element dim list into (B, M, N, K), defaulting B to batch_size.""" + if len(dims) >= 4: + return int(dims[0]), int(dims[1]), int(dims[2]), int(dims[3]) + if len(dims) == 3: + return batch_size, int(dims[0]), int(dims[1]), int(dims[2]) + return batch_size, int(dims[0]), 1, 1 diff --git a/metainfer/tasks/sglang_trace_analyze/orchestrator/fuse_matcher.py b/metainfer/tasks/sglang_trace_analyze/orchestrator/fuse_matcher.py new file mode 100644 index 00000000..58b21ffe --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/orchestrator/fuse_matcher.py @@ -0,0 +1,136 @@ +"""Rule-based fuse pattern matcher. + +Scans the kernel table (ordered by GPU time or timeline order) for known +sequences that indicate a missing fusion opportunity, and reports each +match with a description and estimated saving. + +The catalog is hard-coded — each pattern has a name, the kernel names +that must appear consecutively (or within a short window), and a +suggestion for what the fused replacement would be. +""" + +from __future__ import annotations + +from typing import Any, Dict, List + +# ------------------------------------------------------------------ # +# Fuse pattern catalog +# ------------------------------------------------------------------ # + +FUSE_PATTERNS: List[Dict[str, Any]] = [ + { + "pattern": "rms_norm + gemm", + "kernels": ["rms_norm", "gemm"], + "match_mode": "consecutive", + "suggestion": "Replace separate rms_norm + gemm with fused_rms_norm_gemm (e.g. triton kernel or sglang fused op).", + "estimated_saving_us": 180, + "confidence": "high", + }, + { + "pattern": "silu + mul + gemm", + "kernels": ["silu", "mul", "gemm"], + "match_mode": "consecutive", + "suggestion": "Fuse into silu_and_mul + gemm, or a single fused MoE activation+gemm kernel.", + "estimated_saving_us": 250, + "confidence": "high", + }, + { + "pattern": "add + rms_norm", + "kernels": ["add", "rms_norm"], + "match_mode": "consecutive", + "suggestion": "Fuse residual add + rms_norm into a single kernel to avoid a separate memory round-trip.", + "estimated_saving_us": 120, + "confidence": "medium", + }, + { + "pattern": "quant + gemm", + "kernels": ["quant", "gemm"], + "match_mode": "consecutive", + "suggestion": "Integrate FP8 quantization into the GEMM launch to eliminate a precursor kernel.", + "estimated_saving_us": 200, + "confidence": "medium", + }, + { + "pattern": "nccl_allreduce + gemm (no overlap)", + "kernels": ["ncclAllReduce", "gemm"], + "match_mode": "consecutive", + "suggestion": ( + "AllReduce and gemm are serialized. Try overlapping: issue AllReduce " + "on a separate CUDA stream, or restructure to compute on one output " + "shard while communicating another." + ), + "estimated_saving_us": 300, + "confidence": "medium", + }, +] + + +def match_fuse_patterns( + kernels: List[Dict[str, Any]], +) -> List[Dict[str, Any]]: + """Scan a kernel list for known fuse patterns. + + Args: + kernels: List of kernel entries. Must contain ``kernel_name`` and + preferably be in timeline order. If only duration-ordered, set + ``match_mode`` to ``"unordered"`` for pattern matching. + + Returns: + List of matched patterns, each with ``pattern``, ``kernels``, + ``suggestion``, ``estimated_saving_us``, ``confidence``. + """ + kernel_names = [k.get("kernel_name", "") for k in kernels] + matches = [] + + for pat in FUSE_PATTERNS: + found = _match_consecutive(kernel_names, pat["kernels"]) + if found: + matches.append({ + "pattern": pat["pattern"], + "kernels": found, + "suggestion": pat["suggestion"], + "estimated_saving_us": pat["estimated_saving_us"], + "confidence": pat["confidence"], + }) + + return matches + + +def build_fuse_report( + kernels: List[Dict[str, Any]], + batch_size: int, + stage: str, +) -> Dict[str, Any]: + """Produce the full fuse.json payload.""" + matches = match_fuse_patterns(kernels) + return { + "batch_size": batch_size, + "stage": stage, + "matches": matches, + } + + +def _match_consecutive( + names: List[str], + pattern_kernels: List[str], +) -> List[str]: + """Check if ``pattern_kernels`` appear consecutively (in order) within + ``names``. + + Returns the matched kernel names if found, empty list otherwise. + """ + if len(pattern_kernels) > len(names): + return [] + + patterns_lower = [p.lower() for p in pattern_kernels] + names_lower = [n.lower() for n in names] + + for i in range(len(names_lower) - len(patterns_lower) + 1): + match = True + for j, pat in enumerate(patterns_lower): + if pat not in names_lower[i + j]: + match = False + break + if match: + return names[i: i + len(patterns_lower)] + return [] diff --git a/metainfer/tasks/sglang_trace_analyze/orchestrator/gpu_specs.py b/metainfer/tasks/sglang_trace_analyze/orchestrator/gpu_specs.py new file mode 100644 index 00000000..d88abef7 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/orchestrator/gpu_specs.py @@ -0,0 +1,63 @@ +"""GPU theoretical-peak lookup table. + +Used by ``flops_calculator.py`` to compute MFU: + MFU = actual_TFLOPS / theoretical_peak_TFLOPS. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Dict + + +@dataclass(frozen=True) +class GpuSpec: + """Theoretical peak numbers for one GPU model.""" + + label: str + fp32_tflops: float + tf32_tflops: float + bf16_tflops: float + fp16_tflops: float + int8_tops: float + bandwidth_gb_s: float + + +GPU_SPECS: Dict[str, GpuSpec] = { + "K100": GpuSpec( + label="K100", + fp32_tflops=49, + tf32_tflops=98, + bf16_tflops=192, + fp16_tflops=192, + int8_tops=392, + bandwidth_gb_s=700, + ), + "A100_80G": GpuSpec( + label="A100_80G", + fp32_tflops=19.5, + tf32_tflops=156, + bf16_tflops=312, + fp16_tflops=312, + int8_tops=624, + bandwidth_gb_s=2039, + ), + "H100": GpuSpec( + label="H100", + fp32_tflops=67, + tf32_tflops=989, + bf16_tflops=989, + fp16_tflops=989, + int8_tops=1979, + bandwidth_gb_s=3350, + ), + "B200": GpuSpec( + label="B200", + fp32_tflops=90, + tf32_tflops=2250, + bf16_tflops=2250, + fp16_tflops=2250, + int8_tops=4500, + bandwidth_gb_s=8000, + ), +} diff --git a/metainfer/tasks/sglang_trace_analyze/orchestrator/iteration_record.py b/metainfer/tasks/sglang_trace_analyze/orchestrator/iteration_record.py new file mode 100644 index 00000000..c014831a --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/orchestrator/iteration_record.py @@ -0,0 +1,210 @@ +"""Phase-specific iteration records for sglang_trace_analyze. + +Each phase gets its own dataclass so the schema stays clean — no +``None``-filled optional fields bleeding across phases. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass, field, fields +from typing import Any, Dict + + +def _base_dict(rec, **overrides) -> Dict[str, Any]: + """Serialize *any* iteration record to a dict the WebUI can read. + + Keys: phase (str), status, started_at, ended_at, plus phase-specific + fields from the dataclass. + """ + out: Dict[str, Any] = { + "phase": getattr(rec, "phase", ""), + "status": rec.status, + "started_at": rec.started_at, + "ended_at": rec.ended_at, + } + for f in fields(rec): + if f.name in ("phase", "status", "started_at", "ended_at"): + continue + val = getattr(rec, f.name) + if val is not None: + out[f.name] = val + out.update(overrides) + return out + + +# ------------------------------------------------------------------ # +# MAPPING phase +# ------------------------------------------------------------------ # + +@dataclass +class MappingRecord: + phase: str = "mapping" + status: str = "running" + started_at: float = 0.0 + ended_at: float = 0.0 + batch_size: int | None = None + trace_dir: str | None = None + duration_s: float | None = None + kernel_count: int | None = None + confidence_issues: int = 0 # entries with low confidence after LLM check + error: str | None = None + + def start(self): + self.started_at = time.time() + self.status = "running" + + def done(self, **kw): + self.status = "success" + self.ended_at = time.time() + for k, v in kw.items(): + setattr(self, k, v) + + def fail(self, error: str): + self.status = "failed" + self.ended_at = time.time() + self.error = error + + def to_dict(self) -> Dict[str, Any]: + return _base_dict(self) + + +# ------------------------------------------------------------------ # +# BENCHMARK phase +# ------------------------------------------------------------------ # + +@dataclass +class BenchmarkRecord: + phase: str = "benchmark" + status: str = "running" + started_at: float = 0.0 + ended_at: float = 0.0 + batch_size: int | None = None + trace_dir: str | None = None + duration_s: float | None = None + throughput: float | None = None + latency_p50: float | None = None + error: str | None = None + + def start(self): + self.started_at = time.time() + self.status = "running" + + def done(self, **kw): + self.status = "success" + self.ended_at = time.time() + for k, v in kw.items(): + setattr(self, k, v) + + def fail(self, error: str): + self.status = "failed" + self.ended_at = time.time() + self.error = error + + def to_dict(self) -> Dict[str, Any]: + return _base_dict(self) + + +# ------------------------------------------------------------------ # +# ANALYZE phase +# ------------------------------------------------------------------ # + +@dataclass +class AnalyzeRecord: + phase: str = "analyze" + status: str = "running" + started_at: float = 0.0 + ended_at: float = 0.0 + batch_size: int | None = None + stage: str | None = None # "prefill" | "decode" + kernel_count: int | None = None + top_kernel: str | None = None + top_kernel_pct: float | None = None + mfu_avg: float | None = None + fuse_hits: int = 0 + error: str | None = None + + def start(self): + self.started_at = time.time() + self.status = "running" + + def done(self, **kw): + self.status = "success" + self.ended_at = time.time() + for k, v in kw.items(): + setattr(self, k, v) + + def fail(self, error: str): + self.status = "failed" + self.ended_at = time.time() + self.error = error + + def to_dict(self) -> Dict[str, Any]: + return _base_dict(self) + + +# ------------------------------------------------------------------ # +# HINTS phase +# ------------------------------------------------------------------ # + +@dataclass +class HintsRecord: + phase: str = "hints" + status: str = "running" + started_at: float = 0.0 + ended_at: float = 0.0 + model_used: str | None = None + batch_count: int = 0 + error: str | None = None + + def start(self): + self.started_at = time.time() + self.status = "running" + + def done(self, **kw): + self.status = "success" + self.ended_at = time.time() + for k, v in kw.items(): + setattr(self, k, v) + + def fail(self, error: str): + self.status = "failed" + self.ended_at = time.time() + self.error = error + + def to_dict(self) -> Dict[str, Any]: + return _base_dict(self) + + +# ------------------------------------------------------------------ # +# SUMMARIZE phase +# ------------------------------------------------------------------ # + +@dataclass +class SummarizeRecord: + phase: str = "summarize" + status: str = "running" + started_at: float = 0.0 + ended_at: float = 0.0 + batch_count: int = 0 + best_batch: int | None = None + best_mfu: float | None = None + error: str | None = None + + def start(self): + self.started_at = time.time() + self.status = "running" + + def done(self, **kw): + self.status = "success" + self.ended_at = time.time() + for k, v in kw.items(): + setattr(self, k, v) + + def fail(self, error: str): + self.status = "failed" + self.ended_at = time.time() + self.error = error + + def to_dict(self) -> Dict[str, Any]: + return _base_dict(self) diff --git a/metainfer/tasks/sglang_trace_analyze/orchestrator/orchestrator.py b/metainfer/tasks/sglang_trace_analyze/orchestrator/orchestrator.py new file mode 100644 index 00000000..28f70a9e --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/orchestrator/orchestrator.py @@ -0,0 +1,75 @@ +"""Bootstrap + entry point for the sglang_trace_analyze orchestrator. + +Spawns as a child of the WebUI server per task. Reads requirements, sets +up state_dir / workspace_dir, and runs the linear phase pipeline: + + MAPPING -> BENCHMARK -> ANALYZE -> HINTS -> SUMMARIZE -> done + +Unlike gen_infer_framework, this task has no complex transition table — +just five sequential phases with per-(bs, stage) iterations inside +ANALYZE. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Dict, Optional + +from .pipeline import Pipeline +from .iteration_record import AnalyzeRecord +from metainfer.orchestrator.state import StateStore + + +def run_with_requirements( + requirements_path: Path, + *, + state_dir: Optional[Path] = None, + workspace_dir: Optional[Path] = None, + iter_limit: Optional[int] = None, +) -> int: + """Per-task orchestrator entry point. + + Reads ``requirements.json``, runs the five-phase pipeline to + completion, and exits. + """ + if not requirements_path.exists(): + raise FileNotFoundError(f"requirements file not found: {requirements_path}") + + req: Dict[str, Any] = json.loads( + requirements_path.read_text(encoding="utf-8") + ) + task_id = req.get("task_id", "task") + + # Resolve state_dir + workspace_dir + if state_dir is None or workspace_dir is None: + from metainfer.server import paths as _web_paths + if state_dir is None: + state_dir = _web_paths.task_dir(task_id) + if workspace_dir is None: + workspace_dir = _web_paths.workspace_dir(task_id) + + state_dir.mkdir(parents=True, exist_ok=True) + workspace_dir.mkdir(parents=True, exist_ok=True) + + # Copy requirements into state_dir for self-containment + target_req = state_dir / "requirements.json" + if requirements_path.resolve() != target_req.resolve(): + target_req.write_text( + requirements_path.read_text(encoding="utf-8"), encoding="utf-8" + ) + + store = StateStore(state_dir) + pipe = Pipeline( + req=req, + store=store, + state_dir=state_dir, + workspace_dir=workspace_dir, + ) + + print(f"[metainfer:sglang_trace_analyze] task_id = {task_id}") + print(f"[metainfer:sglang_trace_analyze] state_dir = {state_dir}") + print(f"[metainfer:sglang_trace_analyze] workspace_dir = {workspace_dir}") + + pipe.run() + return 0 diff --git a/metainfer/tasks/sglang_trace_analyze/orchestrator/overlap_detector.py b/metainfer/tasks/sglang_trace_analyze/orchestrator/overlap_detector.py new file mode 100644 index 00000000..e652d66e --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/orchestrator/overlap_detector.py @@ -0,0 +1,115 @@ +"""Detect communication-computation overlap gaps in a torch profiler trace. + +Scans GPU kernel timeline for gaps between consecutive events where the +GPU is idle. On K100, this is lower priority — the detector is kept +simple. +""" + +from __future__ import annotations + +from typing import Any, Dict, List + + +def detect_gaps( + trace_data: Dict[str, Any], + *, + gap_threshold_us: float = 10.0, +) -> List[Dict[str, Any]]: + """Find GPU-idle gaps in the kernel timeline. + + Args: + trace_data: Parsed Chrome trace JSON. + gap_threshold_us: Minimum gap duration (us) to report. + + Returns: + List of gap dicts with ``gap_id``, ``description``, ``gap_us``, + ``affected_kernels``, ``severity``. + """ + trace_events = trace_data.get("traceEvents", []) + if isinstance(trace_data, list): + trace_events = trace_data + + # Collect GPU kernel events with their timestamps + events = [] + for evt in trace_events: + cat = evt.get("cat", "") + dur = evt.get("dur", 0) + ts = evt.get("ts", 0) + if cat == "kernel" and dur > 0: + events.append({ + "name": evt.get("name", ""), + "ts": ts, + "end": ts + dur, + }) + + events.sort(key=lambda e: e["ts"]) + + gaps = [] + gap_id = 0 + for i in range(1, len(events)): + prev_end = events[i - 1]["end"] + curr_start = events[i]["ts"] + gap = curr_start - prev_end + if gap > gap_threshold_us: + gap_id += 1 + severity = "low" + if gap > 100: + severity = "high" + elif gap > 50: + severity = "medium" + + gaps.append({ + "gap_id": gap_id, + "description": ( + f"{events[i - 1]['name']} → {events[i]['name']}: " + f"{gap:.1f}us idle" + ), + "gap_us": round(gap, 1), + "cumulative_gap_us": 0, # filled in by caller + "pct_of_total": 0, # filled in by caller + "affected_kernels": [ + events[i - 1]["name"], + events[i]["name"], + ], + "severity": severity, + }) + + # Compute cumulative stats + total_gap = sum(g["gap_us"] for g in gaps) + total_dur = sum( + (e["end"] - events[0]["ts"]) for e in events[-1:] + ) if events else 0 + + for g in gaps: + g["cumulative_gap_us"] = round(total_gap, 1) + g["pct_of_total"] = round(g["gap_us"] / total_dur * 100, 2) if total_dur > 0 else 0 + + return gaps + + +def build_overlap_report( + trace_data: Dict[str, Any], + batch_size: int, + stage: str, + *, + gap_threshold_us: float = 10.0, +) -> Dict[str, Any]: + """Produce the full overlap.json payload.""" + gaps = detect_gaps(trace_data, gap_threshold_us=gap_threshold_us) + total_gap = sum(g["gap_us"] for g in gaps) + total_dur = sum( + evt.get("dur", 0) for evt in + (trace_data.get("traceEvents", []) or []) + if evt.get("cat") == "kernel" + ) + + return { + "batch_size": batch_size, + "stage": stage, + "gaps": gaps, + "summary": { + "total_gap_us": round(total_gap, 1), + "total_gap_pct": round(total_gap / total_dur * 100, 2) if total_dur > 0 else 0, + "cuda_graph_effective": len(gaps) < 5, + }, + } diff --git a/metainfer/tasks/sglang_trace_analyze/orchestrator/phases.py b/metainfer/tasks/sglang_trace_analyze/orchestrator/phases.py new file mode 100644 index 00000000..a722f6a7 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/orchestrator/phases.py @@ -0,0 +1,55 @@ +"""Phase graph for sglang_trace_analyze. + +Linear pipeline: MAPPING -> BENCHMARK -> ANALYZE -> HINTS -> SUMMARIZE -> done. + +The WebUI state-graph endpoint reads ``terminal_phases`` and +``graph_payload()`` from this module. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +# Ordered list of phases in the pipeline. +PHASES: List[str] = ["mapping", "benchmark", "analyze", "hints", "summarize"] + +# Phases that signal the task is done (whichever is current at exit). +TERMINAL: set[str] = {"done", "failed"} + + +def terminal_phases() -> set[str]: + return TERMINAL + + +def next_phase(current: str) -> str: + """Linear advance. Returns "done" at the end.""" + try: + idx = PHASES.index(current) + if idx + 1 < len(PHASES): + return PHASES[idx + 1] + return "done" + except ValueError: + return "done" + + +def graph_payload( + current: str = "idle", + last_outcome: Optional[str] = None, + last_label: Optional[str] = None, +) -> Dict[str, Any]: + """Return a mermaid-friendly description of the phase graph.""" + nodes = [] + edges = [] + for i, p in enumerate(PHASES): + nodes.append({"id": p, "label": p.upper()}) + if i > 0: + edges.append({"from": PHASES[i - 1], "to": p}) + edges.append({"from": PHASES[-1], "to": "done"}) + nodes.append({"id": "done", "label": "DONE"}) + return { + "nodes": nodes, + "edges": edges, + "current": current, + "last_outcome": last_outcome, + "last_transition_label": last_label, + } diff --git a/metainfer/tasks/sglang_trace_analyze/orchestrator/pipeline.py b/metainfer/tasks/sglang_trace_analyze/orchestrator/pipeline.py new file mode 100644 index 00000000..ee9ca3ef --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/orchestrator/pipeline.py @@ -0,0 +1,995 @@ +"""Pipeline — the sglang_trace_analyze core iteration loop. + +Five-phase linear pipeline: + + MAPPING → BENCHMARK → ANALYZE → HINTS → SUMMARIZE → done + +Each phase may internally iterate (e.g. ANALYZE loops over batch_sizes × +stages). All analysis outputs are written to ``state_dir/analysis/`` as +the authoritative source of truth. +""" + +from __future__ import annotations + +import json +import re +import subprocess +import time +from collections import defaultdict +from pathlib import Path +from typing import Any, Dict, List, Optional + +from metainfer.orchestrator.requirements import req_field, req_field_int +from metainfer.orchestrator.state import StateStore + +from .gpu_specs import GPU_SPECS, GpuSpec +from .iteration_record import ( + AnalyzeRecord, + BenchmarkRecord, + HintsRecord, + MappingRecord, + SummarizeRecord, +) +from .phases import next_phase + + +def _load_json(path: Path, default: Any = None) -> Any: + if not path.exists(): + return default + try: + return json.loads(path.read_text(encoding="utf-8")) + except (ValueError, OSError): + return default + + +def _parse_batch_sizes(raw: str) -> List[int]: + """Parse comma-separated batch sizes, e.g. "1,4,8,16" → [1,4,8,16].""" + return [int(x.strip()) for x in raw.split(",") if x.strip()] + + +def _iter_n(store: StateStore) -> int: + """Next iteration number for timeline ordering. + + Because this pipeline runs phases sequentially with a single iteration + counter (not per-phase counters), we use a simple global counter. + """ + run = store.load_run() + return run.current_iteration + 1 if run else 1 + + +class Pipeline: + """Five-phase profiler-analysis pipeline.""" + + def __init__( + self, + req: Dict[str, Any], + store: StateStore, + state_dir: Path, + workspace_dir: Path, + ): + self.req = req + self.store = store + self.state_dir = state_dir + self.workspace_dir = workspace_dir + self._analysis_dir = state_dir / "analysis" + + # Extract form fields + self.model_path = req_field(req, "model_path", default="") + self.version = req_field(req, "version", default="dev") + self.batch_sizes = _parse_batch_sizes( + req_field(req, "batch_sizes", default="1") + ) + self.mapping_batch_size = req_field_int(req, "mapping_batch_size", default=8) + self.input_len = req_field_int(req, "input_len", default=512) + self.output_len = req_field_int(req, "output_len", default=2000) + self.tp_size = req_field_int(req, "tp_size", default=1) + self.pp_size = req_field_int(req, "pp_size", default=1) + gpu_label = req_field(req, "gpu_model", default="K100") + self.gpu_spec: GpuSpec = GPU_SPECS.get(gpu_label, GPU_SPECS["K100"]) + + # Only decode stage for now + self.stages = ["decode"] # future: ["prefill", "decode"] + + # ------------------------------------------------------------------ # + # Public entry point + # ------------------------------------------------------------------ # + + def run(self) -> None: + """Run the full pipeline.""" + run = self.store.load_run() + phase = run.current_phase or "mapping" + + while phase not in ("done", "failed"): + self.store.update_run(current_phase=phase) + self.store.append_timeline("phase_enter", {"phase": phase}) + + # Skip phases whose outputs already exist (resume / re-run) + if self._phase_is_done(phase): + print(f"[pipeline] phase {phase} output exists, skipping") + self.store.append_timeline("phase_skip", {"phase": phase, "reason": "output exists"}) + phase = next_phase(phase) + continue + + method = getattr(self, f"_run_{phase}", None) + if method is None: + print(f"[pipeline] unknown phase {phase!r}, stopping") + break + + try: + ok = method() + except Exception as exc: + print(f"[pipeline] phase {phase} crashed: {exc}") + self.store.append_timeline("phase_error", {"phase": phase, "error": str(exc)}) + self.store.update_run(finished=True, final_status="failed") + return + + if not ok: + # MAPPING failure is fatal (no traces to analyze). + # BENCHMARK failure is non-fatal: ANALYZE can still use + # MAPPING traces (CUDA Graph OFF) with a note. + if phase == "mapping": + print(f"[pipeline] phase {phase} returned failure, stopping") + self.store.update_run(finished=True, final_status="failed") + return + print(f"[pipeline] phase {phase} returned failure, continuing with available data") + + self.store.append_timeline("phase_exit", {"phase": phase}) + phase = next_phase(phase) + + self.store.update_run(finished=True, final_status="success", + current_phase="done") + self.store.append_timeline("run_done", {"status": "success"}) + + # ================================================================== # + # Phase: MAPPING + # ================================================================== # + + def _run_mapping(self) -> bool: + """Run mapping benchmark (--disable-cuda-graph) then build the + kernel→model-structure mapping from call stacks.""" + print("[pipeline] === MAPPING phase ===") + n = _iter_n(self.store) + rec = MappingRecord(batch_size=self.mapping_batch_size) + rec.start() + self.store.write_iteration(n, rec.to_dict()) + self.store.update_run(current_iteration=n) + + trace_dir = self.workspace_dir / "traces" / "mapping" + + # 1. Generate bench_config.json + bench_config = self._build_bench_config( + batch_sizes=[self.mapping_batch_size], + output_dir=str(self.workspace_dir / "traces"), + ) + config_path = self.state_dir / "bench_config.json" + config_path.write_text(json.dumps(bench_config, indent=2)) + + # 2. Run mapping benchmark + script = Path(__file__).resolve().parent / "run_benchmark.py" + print(f"[pipeline] running mapping benchmark (batch={self.mapping_batch_size})...") + try: + subprocess.run( + [ + "python", str(script), + "--config", str(config_path), + "--mapping-only", + ], + check=True, + timeout=3600, + ) + except subprocess.TimeoutExpired: + rec.fail("mapping benchmark timed out") + self.store.write_iteration(n, rec.to_dict()) + return False + except subprocess.CalledProcessError as e: + rec.fail(f"mapping benchmark exit code {e.returncode}") + self.store.write_iteration(n, rec.to_dict()) + return False + + # 3. Parse trace → build mapping table (rule engine) + # sglang puts traces inside a timestamp subdirectory + decode_trace_dir = trace_dir + if not decode_trace_dir.exists(): + # try globbing for timestamp subdirs + ts_dirs = sorted(trace_dir.parent.glob( + trace_dir.name + "/*" if trace_dir.name else "*/" + )) if trace_dir.parent.exists() else [] + if not ts_dirs: + # fall back: find any trace files + candidates = list(trace_dir.parent.rglob("*.trace.json.gz")) if trace_dir.parent.exists() else [] + if not candidates: + rec.fail("no trace files found after mapping benchmark") + self.store.write_iteration(n, rec.to_dict()) + return False + trace_path = candidates[0] + else: + decode_trace_dir = ts_dirs[0] + traces = sorted(decode_trace_dir.glob("*DECODE*.trace.json.gz")) + if not traces: + rec.fail("no decode traces in " + str(decode_trace_dir)) + self.store.write_iteration(n, rec.to_dict()) + return False + trace_path = traces[0] + else: + traces = sorted(decode_trace_dir.glob("*DECODE*.trace.json.gz")) + if not traces: + traces = sorted(decode_trace_dir.rglob("*DECODE*.trace.json.gz")) + if not traces: + rec.fail("no decode traces found in " + str(decode_trace_dir)) + self.store.write_iteration(n, rec.to_dict()) + return False + trace_path = traces[0] + + print(f"[pipeline] parsing trace: {trace_path}") + mapping_entries = self._build_mapping(trace_path) + if not mapping_entries: + rec.fail("mapping produced zero entries — trace may be empty or format unsupported") + self.store.write_iteration(n, rec.to_dict()) + return False + + # 4. LLM sanity check (placeholder — calls sub-agent when available) + mapping_entries = self._llm_mapping_sanity_check(mapping_entries) + + # 5. Write mapping.json + self._analysis_dir.mkdir(parents=True, exist_ok=True) + mapping_file = self._analysis_dir / "mapping.json" + mapping_file.write_text(json.dumps({ + "model": self.model_path, + "gpu": self.gpu_spec.label, + "mapping_batch_size": self.mapping_batch_size, + "entries": mapping_entries, + }, indent=2)) + + confidence_issues = sum( + 1 for e in mapping_entries + if e.get("confidence", "high") == "low" + ) + + rec.done( + trace_dir=str(trace_dir), + kernel_count=len(mapping_entries), + confidence_issues=confidence_issues, + ) + self.store.write_iteration(n, rec.to_dict()) + return True + + # ------------------------------------------------------------------ # + # Phase: BENCHMARK + # ------------------------------------------------------------------ # + + def _run_benchmark(self) -> bool: + """Run formal benchmarks (CUDA Graph ON) for each batch size.""" + print("[pipeline] === BENCHMARK phase ===") + + bench_config = self._build_bench_config( + batch_sizes=self.batch_sizes, + output_dir=str(self.workspace_dir / "traces"), + ) + config_path = self.state_dir / "bench_config.json" + config_path.write_text(json.dumps(bench_config, indent=2)) + + script = Path(__file__).resolve().parent / "run_benchmark.py" + + all_ok = True + for bs in self.batch_sizes: + n = _iter_n(self.store) + rec = BenchmarkRecord(batch_size=bs) + rec.start() + self.store.write_iteration(n, rec.to_dict()) + self.store.update_run(current_iteration=n) + + trace_dir = self.workspace_dir / "traces" / f"bs_{bs}" + print(f"[pipeline] batch_size={bs}") + + try: + subprocess.run( + [ + "python", str(script), + "--config", str(config_path), + "--formal-only", + "--single-batch", str(bs), + ], + check=True, + timeout=3600, + ) + except subprocess.TimeoutExpired: + rec.fail("timed out") + self.store.write_iteration(n, rec.to_dict()) + all_ok = False + continue + except subprocess.CalledProcessError as e: + rec.fail(f"exit code {e.returncode}") + self.store.write_iteration(n, rec.to_dict()) + all_ok = False + continue + + # Extract throughput/latency from sglang output log if available + rec.done(trace_dir=str(trace_dir)) + self.store.write_iteration(n, rec.to_dict()) + + # Continue if any traces exist (mapping or formal) + has_formal = any( + (self.workspace_dir / "traces" / f"bs_{bs}").exists() + for bs in self.batch_sizes + ) + has_mapping = (self.workspace_dir / "traces" / "mapping").exists() + return all_ok or has_formal or has_mapping + + # ================================================================== # + # Phase: ANALYZE + # ================================================================== # + + def _run_analyze(self) -> bool: + """Analyze each (batch_size, stage) pair that has a trace.""" + print("[pipeline] === ANALYZE phase ===") + + mapping = _load_json(self._analysis_dir / "mapping.json", {}) + mapping_entries = mapping.get("entries", []) + if not mapping_entries: + print("[pipeline] WARNING: no mapping entries — analysis may be incomplete") + + any_ok = False + for bs in self.batch_sizes: + for stage in self.stages: + # Priority: formal CUDA Graph ON traces > mapping traces + trace_dir = self._find_trace_dir(bs, stage) + if trace_dir is None: + print(f"[pipeline] skipping bs_{bs}/{stage} — no trace dir") + continue + + traces = sorted(trace_dir.glob("*DECODE*.trace.json.gz")) + if not traces: + print(f"[pipeline] skipping bs_{bs}/{stage} — no trace files") + continue + + n = _iter_n(self.store) + rec = AnalyzeRecord(batch_size=bs, stage=stage) + rec.start() + self.store.write_iteration(n, rec.to_dict()) + self.store.update_run(current_iteration=n) + + print(f"[pipeline] analyzing bs_{bs}/{stage} ({traces[0].name})") + try: + result = self._analyze_one( + traces[0], mapping_entries, bs, stage + ) + except Exception as exc: + rec.fail(str(exc)) + self.store.write_iteration(n, rec.to_dict()) + continue + + # Write output files + out_dir = self._analysis_dir / "batches" / f"bs_{bs}" / stage + out_dir.mkdir(parents=True, exist_ok=True) + + (out_dir / "kernel_table.json").write_text( + json.dumps(result["kernel_table"], indent=2)) + (out_dir / "overlap.json").write_text( + json.dumps(result["overlap"], indent=2)) + (out_dir / "fuse.json").write_text( + json.dumps(result["fuse"], indent=2)) + + mfu_vals = [ + k.get("mfu", 0) for k in result["kernel_table"].get("kernels", []) + if k.get("mfu") is not None + ] + top_kernels = result["kernel_table"].get("kernels", []) + rec.done( + kernel_count=len(top_kernels), + top_kernel=top_kernels[0]["kernel_name"] if top_kernels else None, + top_kernel_pct=top_kernels[0]["time_pct"] if top_kernels else None, + mfu_avg=round(sum(mfu_vals) / len(mfu_vals), 1) if mfu_vals else None, + fuse_hits=len(result["fuse"].get("matches", [])), + ) + self.store.write_iteration(n, rec.to_dict()) + any_ok = True + + return any_ok + + # ================================================================== # + # Phase: HINTS + # ================================================================== # + + def _run_hints(self) -> bool: + """Generate LLM optimization hints from all analysis results.""" + print("[pipeline] === HINTS phase ===") + n = _iter_n(self.store) + rec = HintsRecord( + model_used=self.model_path, + batch_count=len(self.batch_sizes), + ) + rec.start() + self.store.write_iteration(n, rec.to_dict()) + self.store.update_run(current_iteration=n) + + # Collect full kernel tables + summaries from all analyzed batches + kernel_summaries = [] + overlap_summaries = [] + fuse_summaries = [] + + for bs in self.batch_sizes: + for stage in self.stages: + out_dir = self._analysis_dir / "batches" / f"bs_{bs}" / stage + kt = _load_json(out_dir / "kernel_table.json") + ov = _load_json(out_dir / "overlap.json") + fu = _load_json(out_dir / "fuse.json") + if kt: + kernel_summaries.append({ + "batch_size": bs, "stage": stage, + "all_kernels": kt.get("kernels", []), + }) + if ov: + overlap_summaries.append(ov) + if fu: + fuse_summaries.append(fu) + + if not kernel_summaries: + print("[pipeline] no kernel tables — skipping hints") + rec.fail("no analysis data available") + self.store.write_iteration(n, rec.to_dict()) + return True # not fatal — hints are optional + + # Generate hints (placeholder — real impl calls LLM sub-agent) + hints = self._llm_generate_hints( + kernel_summaries, overlap_summaries, fuse_summaries + ) + self._analysis_dir.mkdir(parents=True, exist_ok=True) + (self._analysis_dir / "hints.json").write_text( + json.dumps(hints, indent=2)) + + rec.done() + self.store.write_iteration(n, rec.to_dict()) + return True + + # ================================================================== # + # Phase: SUMMARIZE + # ================================================================== # + + def _run_summarize(self) -> bool: + """Aggregate cross-batch summary.""" + print("[pipeline] === SUMMARIZE phase ===") + n = _iter_n(self.store) + rec = SummarizeRecord(batch_count=len(self.batch_sizes)) + rec.start() + self.store.write_iteration(n, rec.to_dict()) + self.store.update_run(current_iteration=n) + + batch_summaries = [] + best_batch = None + best_mfu = None + + for bs in self.batch_sizes: + for stage in ["decode"]: + out_dir = self._analysis_dir / "batches" / f"bs_{bs}" / stage + kt = _load_json(out_dir / "kernel_table.json") + if not kt: + batch_summaries.append({ + "batch_size": bs, "stage": stage, + "status": "missing", + }) + continue + + kernels = kt.get("kernels", []) or [] + mfu_vals = [k.get("mfu", 0) for k in kernels if k.get("mfu")] + avg_mfu = round(sum(mfu_vals) / len(mfu_vals), 1) if mfu_vals else None + top = kernels[0] if kernels else {} + + info = { + "batch_size": bs, + "stage": stage, + "top_kernel": top.get("kernel_name"), + "top_kernel_pct": top.get("time_pct"), + "mfu_avg": avg_mfu, + "kernel_count": len(kernels), + } + batch_summaries.append(info) + + if avg_mfu is not None and (best_mfu is None or avg_mfu > best_mfu): + best_mfu = avg_mfu + best_batch = bs + + self._analysis_dir.mkdir(parents=True, exist_ok=True) + (self._analysis_dir / "summary.json").write_text(json.dumps({ + "model": self.model_path, + "gpu": self.gpu_spec.label, + "batches": batch_summaries, + }, indent=2)) + + rec.done(best_batch=best_batch, best_mfu=best_mfu) + self.store.write_iteration(n, rec.to_dict()) + return True + + # ================================================================== # + # Phase skip detection (resume / re-run) + # ================================================================== # + + def _phase_is_done(self, phase: str) -> bool: + """Return True if the phase's expected outputs already exist.""" + if phase == "mapping": + return (self._analysis_dir / "mapping.json").exists() + if phase == "benchmark": + # Check that at least one batch_size trace dir exists + for bs in self.batch_sizes: + trace_dir = self.workspace_dir / "traces" / f"bs_{bs}" + if trace_dir.exists(): + return True + return False + if phase == "analyze": + for bs in self.batch_sizes: + for stage in self.stages: + if not (self._analysis_dir / "batches" / f"bs_{bs}" / stage / "kernel_table.json").exists(): + return False + return True + if phase == "hints": + return (self._analysis_dir / "hints.json").exists() + if phase == "summarize": + return (self._analysis_dir / "summary.json").exists() + return False + + # ================================================================== # + # Trace discovery + # ================================================================== # + + def _find_trace_dir(self, bs: int, stage: str) -> Optional[Path]: + """Find the best trace directory for a (batch_size, stage) pair. + + Priority: formal traces (CUDA Graph ON, under ``bs_/``) > + mapping traces (CUDA Graph OFF, under ``mapping/``). + + sglang ``--profile-by-stage`` saves traces inside a timestamp + subdirectory, so we look there first. + """ + def _find_in(base: Path) -> Optional[Path]: + if not base.exists(): + return None + # Direct: bs_8/decode/*.trace.json.gz + direct = base / stage + if direct.exists(): + traces = list(direct.glob("*DECODE*.trace.json.gz")) + if traces: + return direct + # Timestamp subdir: bs_8//*.trace.json.gz + ts_dirs = sorted([d for d in base.iterdir() if d.is_dir()]) + for ts in ts_dirs: + traces = list(ts.glob("*DECODE*.trace.json.gz")) + if traces: + return ts + return None + + # 1. Formal traces + formal_base = self.workspace_dir / "traces" / f"bs_{bs}" + found = _find_in(formal_base) + if found: + return found + + # 2. Mapping traces + map_base = self.workspace_dir / "traces" / "mapping" + found = _find_in(map_base) + if found: + return found + + return None + + # ================================================================== # + # Helpers + # ================================================================== # + + def _build_bench_config( + self, batch_sizes: List[int], output_dir: str, + ) -> Dict[str, Any]: + return { + "model_path": self.model_path, + "version": self.version, + "batch_sizes": batch_sizes, + "mapping_batch_size": self.mapping_batch_size, + "input_len": self.input_len, + "output_len": self.output_len, + "tp_size": self.tp_size, + "pp_size": self.pp_size, + "output_dir": output_dir, + "profile_start_step": 5, + "profile_steps": 5, + } + + def _is_formal_trace(self, trace_path) -> bool: + """Return True if this is a formal (CUDA Graph ON) trace.""" + return "_graph_" in str(trace_path) and "_nograph_" not in str(trace_path) + + def _merge_mapping_tflops( + self, result_kernels: list, bs: int, stage: str + ) -> None: + """Enrich formal trace kernel entries with TFLOPS/MFU/bound from + the mapping trace, which has per-kernel Input Dims. + + Matches kernels by name and overwrites tflops_actual, mfu, + bound, bandwidth_gb_s, and input_dims from the mapping trace. + """ + mapping_trace_dir = self._find_mapping_trace_dir() + if mapping_trace_dir is None: + print("[pipeline] no mapping trace to enrich TFLOPS data") + return + + traces = sorted(mapping_trace_dir.glob("*DECODE*.trace.json.gz")) + if not traces: + return + + from .trace_parser import parse_trace, aggregate_kernels + from .flops_calculator import calculate_mfu + + print(f"[pipeline] enriching TFLOPS from mapping trace") + map_data = parse_trace(str(traces[0])) + map_kernels = aggregate_kernels(map_data) + + # Build CPU op correlation for mapping trace too + events = map_data.get("traceEvents", []) + cpu_ops_by_corr = defaultdict(lambda: []) + kernel_by_corr = defaultdict(lambda: []) + for e in events: + cat = e.get("cat", "") + corr = (e.get("args") or {}).get("External id" if cat == "cpu_op" else "correlation") + if cat == "cpu_op" and corr: + cpu_ops_by_corr[corr].append(e.get("name", "")) + elif cat == "kernel" and corr: + kernel_by_corr[corr].append(e.get("name", "")) + + # Classify + calculate MFU + from .structure_mapper import _map_one as map_one + map_entries = [] + for k in map_kernels: + name = k["kernel_name"] + cpu_ops = set() + for corr, gpu_names in kernel_by_corr.items(): + if name in gpu_names: + for cn in cpu_ops_by_corr.get(corr, []): + cpu_ops.add(cn) + mapped = map_one(name, k.get("call_stack", ""), {}, list(cpu_ops)) + k.update(mapped) + map_entries.append(k) + + map_entries = calculate_mfu(map_entries, self.gpu_spec, batch_size=bs, dtype="bf16") + + # Build lookup by kernel name + map_lookup = {k["kernel_name"]: k for k in map_entries} + + enriched = 0 + for k in result_kernels: + name = k["kernel_name"] + if name in map_lookup: + src = map_lookup[name] + if src.get("tflops_actual") is not None: + k["tflops_actual"] = src["tflops_actual"] + k["mfu"] = src["mfu"] + k["bound"] = src["bound"] + k["bandwidth_gb_s"] = src["bandwidth_gb_s"] + k["input_dims"] = src.get("input_dims", []) + k["flops_per_invocation"] = src.get("flops_per_invocation", 0) + enriched += 1 + print(f"[pipeline] enriched {enriched}/{len(result_kernels)} kernels with TFLOPS from mapping trace") + + def _find_mapping_trace_dir(self) -> Optional[Path]: + """Find the mapping trace directory (CUDA Graph OFF).""" + map_base = self.workspace_dir / "traces" / "mapping" + if not map_base.exists(): + return None + # Check for timestamp subdirs first + ts_dirs = sorted([d for d in map_base.iterdir() if d.is_dir()]) + for ts in ts_dirs: + traces = list(ts.glob("*DECODE*.trace.json.gz")) + if traces: + return ts + # Direct + traces = list(map_base.glob("*DECODE*.trace.json.gz")) + if traces: + return map_base + return None + + def _build_mapping(self, trace_path: Path) -> List[Dict[str, Any]]: + """Parse a trace file and build kernel→model-structure mapping + using trace_parser + structure_mapper with CPU op correlation.""" + from .trace_parser import parse_trace, aggregate_kernels + from .structure_mapper import _map_one as map_one + + print(f"[pipeline] parsing trace for mapping: {trace_path}") + trace_data = parse_trace(str(trace_path)) + + kernels = aggregate_kernels(trace_data) + + # Build CPU op correlation + events = trace_data.get("traceEvents", []) + cpu_ops_by_corr = defaultdict(lambda: []) + kernel_by_corr = defaultdict(lambda: []) + for e in events: + cat = e.get("cat", "") + corr = (e.get("args") or {}).get( + "External id" if cat == "cpu_op" else "correlation" + ) + if cat == "cpu_op" and corr: + cpu_ops_by_corr[corr].append(e.get("name", "")) + elif cat == "kernel" and corr: + kernel_by_corr[corr].append(e.get("name", "")) + + entries = [] + seen = set() + for k in kernels: + name = k["kernel_name"] + if name in seen: + continue + seen.add(name) + cpu_ops = set() + for corr, gpu_names in kernel_by_corr.items(): + if name in gpu_names: + for cn in cpu_ops_by_corr.get(corr, []): + cpu_ops.add(cn) + entry = map_one(name, k.get("call_stack", ""), {}, list(cpu_ops)) + entries.append(entry) + + print(f"[pipeline] mapping built: {len(entries)} unique kernels") + return entries + + def _llm_mapping_sanity_check( + self, entries: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: + """Run LLM sanity check on mapping entries. + + Placeholder — real impl calls SubAgentManager. + """ + # TODO: wire SubAgentManager + for e in entries: + e.setdefault("confidence", "high") + return entries + + def _analyze_one( + self, + trace_path: Path, + mapping_entries: List[Dict[str, Any]], + bs: int, + stage: str, + ) -> Dict[str, Any]: + """Analyze a single trace file — uses trace_parser, structure_mapper, + flops_calculator, overlap_detector, fuse_matcher.""" + from .trace_parser import parse_trace, aggregate_kernels + from .structure_mapper import _map_one as map_one + from .flops_calculator import calculate_mfu + from .overlap_detector import build_overlap_report + from .fuse_matcher import build_fuse_report + + print(f"[pipeline] loading trace: {trace_path}") + trace_data = parse_trace(str(trace_path)) + + # Aggregate kernels + kernels = aggregate_kernels(trace_data) + total_dur = sum(k["total_dur_us"] for k in kernels) / 1e6 + + # Build CPU op correlation + events = trace_data.get("traceEvents", []) + cpu_ops_by_corr = defaultdict(lambda: []) + kernel_by_corr = defaultdict(lambda: []) + for e in events: + cat = e.get("cat", "") + corr = (e.get("args") or {}).get( + "External id" if cat == "cpu_op" else "correlation" + ) + if cat == "cpu_op" and corr: + cpu_ops_by_corr[corr].append(e.get("name", "")) + elif cat == "kernel" and corr: + kernel_by_corr[corr].append(e.get("name", "")) + + # Map each kernel using structure_mapper + cpu_ops + result_kernels = [] + for k in kernels: + name = k["kernel_name"] + # Collect correlated CPU ops + cpu_ops = set() + for corr, gpu_names in kernel_by_corr.items(): + if name in gpu_names: + for cn in cpu_ops_by_corr.get(corr, []): + cpu_ops.add(cn) + + # Use structure_mapper for op_type/category/layer + mapped = map_one(name, k.get("call_stack", ""), {}, list(cpu_ops)) + pct = k["total_dur_us"] / (total_dur * 1e6) * 100 + + entry = { + "rank": len(result_kernels) + 1, + "kernel_name": name, + "category": mapped["category"], + "op_type": mapped["op_type"], + "model_layer": mapped["model_layer"], + "confidence": mapped["confidence"], + "total_dur_us": k["total_dur_us"], + "time_pct": round(pct, 2), + "count": k["count"], + "avg_dur_us": round(k["total_dur_us"] / k["count"], 2) if k["count"] else 0, + "input_dims": k.get("input_dims", []), + "tflops_theoretical": self.gpu_spec.bf16_tflops, + "bandwidth_theoretical": self.gpu_spec.bandwidth_gb_s, + } + result_kernels.append(entry) + + # Calculate TFLOPS/MFU/bound using flops_calculator + result_kernels = calculate_mfu( + result_kernels, self.gpu_spec, batch_size=bs, dtype="bf16" + ) + + kernel_table = { + "model": self.model_path, + "gpu": self.gpu_spec.label, + "batch_size": bs, + "stage": stage, + "total_gpu_time_s": round(total_dur, 2), + "unique_kernels": len(result_kernels), + "kernels": result_kernels, + } + + # Enrich formal traces with TFLOPS from mapping trace + if self._is_formal_trace(trace_path): + self._merge_mapping_tflops(result_kernels, bs, stage) + + overlap = build_overlap_report(trace_data, bs, stage) + # Detect CUDA Graph from trace filename: _graph_ = formal, _nograph_ = mapping + trace_name = str(trace_path) + if "_nograph_" in trace_name: + overlap["summary"]["cuda_graph_effective"] = False + elif "_graph_" in trace_name: + overlap["summary"]["cuda_graph_effective"] = True + fuse = build_fuse_report(result_kernels, bs, stage) + + return {"kernel_table": kernel_table, "overlap": overlap, "fuse": fuse} + + def _llm_generate_hints( + self, + kernel_summaries: list, + overlap_summaries: list, + fuse_summaries: list, + ) -> Dict[str, Any]: + """Generate optimization hints from analysis data. + + Uses rule-based analysis of kernel tables, overlap, and fuse results + to produce actionable optimization suggestions. + """ + suggestions = [] + surprises = [] + + # Collect all kernels across batches/stages + all_kernels = [] + for ks in kernel_summaries: + for k in (ks.get("all_kernels") or []): + all_kernels.append(k) + + if not all_kernels: + return { + "bottleneck": {"kernel_or_pattern": "unknown", "reason": "no data", "impact_pct": 0}, + "suggestions": [], "surprises": [], + "status": "generated", + } + + top = all_kernels[0] if all_kernels else {} + top_name = top.get("kernel_name", "unknown") + top_cat = top.get("category", "Other") + top_pct = top.get("time_pct", 0) + + # Categorize kernels + cats = {} + for k in all_kernels: + c = k.get("category", "Other") + cats[c] = cats.get(c, 0) + (k.get("time_pct", 0) or 0) + + # 1. Dominant kernel analysis + if top_pct > 50: + suggestions.append({ + "title": f"Replace or optimize {top_cat} kernel", + "what_to_change": f"The \"{top_name[:40]}\" kernel dominates at {top_pct:.0f}% GPU time. Profile with Nsight Compute to identify micro-architectural bottlenecks, or replace with a vendor-optimized implementation.", + "why": f"Single kernel consuming >50% of GPU time is the highest-ROI optimization target.", + "estimated_saving_pct": round(top_pct * 0.3), + "difficulty": "high", + "category": "kernel_replace", + }) + elif top_pct > 20: + suggestions.append({ + "title": f"Profile {top_cat} kernel with Nsight", + "what_to_change": f"\"{top_name[:40]}\" at {top_pct:.0f}%. Use Nsight Compute to check occupancy, memory coalescing, and register pressure.", + "why": "Top kernel is a clear bottleneck. Micro-architectural optimization may yield 10-30% improvement.", + "estimated_saving_pct": round(top_pct * 0.2), + "difficulty": "medium", + "category": "kernel_replace", + }) + + # 2. Category-specific suggestions + reduce_pct = cats.get("Reduce", 0) + if reduce_pct > 30: + suggestions.append({ + "title": "Reduce TP allreduce overhead", + "what_to_change": "Custom allreduce consumes {:.0f}% GPU time. Try: (1) overlap communication with computation using separate CUDA streams, (2) reduce TP degree if memory permits, or (3) enable CUDA Graph to amortize launch overhead.".format(reduce_pct), + "why": "TP communication is the dominant cost. Even 10% reduction saves significant time.", + "estimated_saving_pct": round(reduce_pct * 0.25), + "difficulty": "medium", + "category": "overlap", + }) + + gemm_pct = cats.get("GEMM", 0) + if gemm_pct > 20: + suggestions.append({ + "title": "Quantize GEMMs to FP8 or INT8", + "what_to_change": "GEMM kernels consume {:.0f}% GPU time. Explore FP8 (w8a8) quantization for attention projections and FFN layers to double throughput.".format(gemm_pct), + "why": "GEMM is compute-heavy and benefits most from reduced precision.", + "estimated_saving_pct": round(gemm_pct * 0.4), + "difficulty": "medium", + "category": "config_tune", + }) + + element_pct = cats.get("ElementWise", 0) + if element_pct > 15: + suggestions.append({ + "title": "Fuse element-wise operations", + "what_to_change": f"Element-wise kernels consume {element_pct:.0f}% GPU time. These are memory-bound — fuse consecutive element-wise ops (add, mul, silu, norm) into single kernels to reduce memory traffic.", + "why": "Memory-bound element-wise ops benefit most from fusion, eliminating intermediate reads/writes.", + "estimated_saving_pct": round(element_pct * 0.4), + "difficulty": "low", + "category": "fuse", + }) + + # 3. CUDA Graph check (from overlap data) + any_cuda_graph = any( + s.get("summary", {}).get("cuda_graph_effective", False) + for s in overlap_summaries + ) + if not any_cuda_graph: + suggestions.append({ + "title": "Enable CUDA Graph for decode", + "what_to_change": "CUDA Graph is not active. Enable --cuda-graph-bs to capture and replay the decode graph. On K100 with DeepSeek V4, this typically yields 3-5x throughput improvement.", + "why": "Decode is launch-bound. CUDA Graph eliminates per-step kernel launch overhead.", + "estimated_saving_pct": 70, + "difficulty": "low", + "category": "config_tune", + }) + + # 4. Surprises + if reduce_pct < 5 and "NCCL" not in cats: + surprises.append("TP allreduce overhead is unexpectedly low — verify communication is actually happening (check TP degree).") + if gemm_pct > 50: + surprises.append("GEMM dominates at >50% — unexpected for a decode workload. Check if attention is correctly fused.") + + bottleneck = { + "kernel_or_pattern": top_name[:80] if top_name else "unknown", + "reason": f"Largest single consumer of GPU time at {top_pct:.1f}% (category: {top_cat})", + "impact_pct": round(top_pct), + } + + return { + "bottleneck": bottleneck, + "suggestions": suggestions[:5], + "surprises": surprises, + "status": "generated", + } + + +# ------------------------------------------------------------------ # +# Module-level kernel classifier +# ------------------------------------------------------------------ # + +def _classify_kernel(name: str) -> tuple: + """Classify a GPU kernel name into (op_type, category).""" + n = name.lower() + if n.startswith("cijk_"): + return ("GEMM", "CK-GEMM") + if "flash_fwd" in n or "flash_attn" in n: + return ("Attention", "MLA") + if "fused_moe" in n: + return ("MoE", "MoE") + if "nccl" in n: + return ("NCCL", "NCCL-AllGather") + if "allreduce" in n: + return ("NCCL", "NCCL-AllReduce") + if "reduce_kernel" in n: + return ("Reduce", "CustomAllReduce") + if "rms_norm" in n or "rmsnorm" in n: + return ("Norm", "RMSNorm") + if "elementwise" in n: + return ("ElementWise", "ElementWise") + if "gather" in n or "topk" in n: + return ("Memory", "Gather") + if "copy" in n or "memcpy" in n: + return ("Memory", "Copy") + if "vectorized" in n: + return ("ElementWise", "Vectorized") + return ("Other", "Other") diff --git a/metainfer/tasks/sglang_trace_analyze/orchestrator/plugin.py b/metainfer/tasks/sglang_trace_analyze/orchestrator/plugin.py new file mode 100644 index 00000000..d3d14812 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/orchestrator/plugin.py @@ -0,0 +1,10 @@ +"""TaskPlugin descriptor for sglang_trace_analyze.""" + +from metainfer.orchestrator.tasks.base import TaskPlugin + +PLUGIN = TaskPlugin( + task_type="sglang_trace_analyze", + cli_module="metainfer.tasks.sglang_trace_analyze.orchestrator.cli", + phases_module="metainfer.tasks.sglang_trace_analyze.orchestrator.phases", + diagnostic_globs=("*",), +) diff --git a/metainfer/tasks/sglang_trace_analyze/orchestrator/prompts.py b/metainfer/tasks/sglang_trace_analyze/orchestrator/prompts.py new file mode 100644 index 00000000..2b4994b4 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/orchestrator/prompts.py @@ -0,0 +1,117 @@ +"""LLM prompts for sglang_trace_analyze. + +Two prompt families: +1. **mapping_sanity_check** — validate the kernel-to-model-structure mapping. +2. **optimization_hints** — generate actionable optimization suggestions from + the full analysis (kernel tables + overlap + fuse results). +""" + +from __future__ import annotations + + +def mapping_sanity_check_prompt( + mapping_json: str, + model_config_json: str, + gpu_label: str, +) -> str: + """Prompt for LLM to sanity-check a kernel → model-structure mapping.""" + return f"""You are a GPU inference optimization expert. Review the following +kernel-to-model-structure mapping that was auto-generated from a torch +profiler trace's call stacks. + +## Model config.json +```json +{model_config_json} +``` + +## Auto-generated mapping (excerpt — full file too large, this is the first +200 entries sorted by GPU time) +```json +{mapping_json} +``` + +## GPU +{gpu_label} + +## Tasks +1. For each mapping entry, rate its confidence: "high" (call stack clearly + points to a known layer/op), "medium" (plausible but ambiguous), or + "low" (likely wrong — kernel name and call stack don't match expected + pattern). If you're uncertain about a model architecture detail, search + the web for the model's architecture documentation before rating. +2. Flag any kernel that appears to be mapped to the wrong layer type + (e.g. a MoE kernel mapped to a dense layer, or an attention kernel + mapped to an FFN layer). +3. Flag missing mappings — kernel names that appear in the trace but have + no clear model-layer assignment. +4. Return a JSON object with this schema: + {{ + "entries": [ + {{ + "kernel_name": "...", + "confidence": "high|medium|low", + "issues": ["..."] // empty list if none + }} + ], + "summary": {{ + "high_count": N, + "medium_count": N, + "low_count": N, + "overall_assessment": "..." + }} + }} +""" + + +def optimization_hints_prompt( + kernel_tables_summary: str, + overlap_summary: str, + fuse_summary: str, + gpu_label: str, + model_name: str, +) -> str: + """Prompt for LLM to generate optimization hints from analysis results.""" + return f"""You are a GPU inference optimization expert. Review the profiling +analysis below and generate actionable optimization suggestions. + +## Model +{model_name} + +## GPU +{gpu_label} + +## Kernel Hotspot Summary (top kernels by GPU time across all batch sizes) +{kernel_tables_summary} + +## Overlap Analysis +{overlap_summary} + +## Fuse Pattern Matches +{fuse_summary} + +## Tasks +1. Identify the single biggest bottleneck and explain why it dominates. +2. List 3-5 concrete optimization directions, ordered by estimated impact. + For each: what to change, why it helps, and estimated saving (%). +3. Note any surprising or counter-intuitive findings (e.g. a kernel that + should be fast but is unexpectedly slow). +4. Return a JSON object with this schema: + {{ + "bottleneck": {{ + "kernel_or_pattern": "...", + "reason": "...", + "impact_pct": N + }}, + "suggestions": [ + {{ + "title": "...", + "what_to_change": "...", + "why": "...", + "estimated_saving_pct": N, + "difficulty": "low|medium|high", + "category": "fuse|overlap|kernel_replace|config_tune|other" + }} + ], + "surprises": ["..."] + }} +""" diff --git a/metainfer/tasks/sglang_trace_analyze/orchestrator/run_benchmark.py b/metainfer/tasks/sglang_trace_analyze/orchestrator/run_benchmark.py new file mode 100644 index 00000000..9f90d079 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/orchestrator/run_benchmark.py @@ -0,0 +1,205 @@ +#!/usr/bin/env python3 +"""Wrapper script for sglang.bench_one_batch_server. + +Called by the orchestrator pipeline in two modes: + + # Mapping run — one batch size, --disable-cuda-graph + python run_benchmark.py --config bench_config.json --mapping-only + + # Formal runs — one or all batch sizes, CUDA Graph ON + python run_benchmark.py --config bench_config.json --formal-only [--single-batch N] + +Environment variables required for K100/HIP are set inside this script +so callers don't need to source them externally. +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +from pathlib import Path +from typing import Any, Dict, List + +# ── K100 / HIP environment — must match /workspace/sglang/scripts/run_traces.sh ── + +_K100_ENV = { + "SGL_CHUNKED_PREFIX_CACHE_THRESHOLD": "0", + "SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT": "1200", + "GLIBC_TUNABLES": "glibc.rtld.optional_static_tls=0x40000", + "SGLANG_SET_CPU_AFFINITY": "1", + "HIP_KERNEL_BATCH_CEILING": "100", + "GPU_MAX_HW_QUEUES": "3", + # "HIP_GRAPH_ACCUMULATE_DISPATCH": "0", # torchprof needs this + "HIP_H2D_DISABLE_COPY_BUFFER": "0", + "HIP_D2H_DISABLE_COPY_BUFFER": "0", + "HIP_H2D_DIRECT_COPY_THRESHOLD": "32768", + "HIP_H2D_HSAAPI_COPY_THRESHOLD": "32768", + "HIP_D2H_DIRECT_COPY_THRESHOLD": "512", + "HIP_D2H_HSAAPI_COPY_THRESHOLD": "512", + "USE_DCU_CUSTOM_ALLREDUCE": "1", + "HIP_KERNEL_EVENT_SYSTENFENCE": "1", + "SGLANG_USE_FP8_W8A8_MOE": "0", + "SGLANG_USE_LIGHTOP": "0", + "SGLANG_ROCM_USE_AITER_MOE": "0", + "SGLANG_OPT_USE_FUSED_HASH_TOPK": "false", + "SGLANG_OPT_SWIGLU_CLAMP_FUSION": "false", + "SGLANG_TOPK_TRANSFORM_512_TORCH": "false", + "SGLANG_OPT_USE_JIT_KERNEL_FUSED_TOPK": "false", + "SGLANG_NSA_FUSE_TOPK": "false", + "SGLANG_JIT_DEEPGEMM_PRECOMPILE": "0", + "SGLANG_APPLY_CONFIG_BACKUP": "none", + "SGLANG_DSV4_MODE": "2604", + "SGLANG_OPT_BF16_FP32_GEMM_ALGO": "torch", + "SGLANG_OPT_USE_HIP_PAGED_MQA_LOGITS": "1", + "SGLANG_OPT_USE_HIP_MHC_PRE": "1", + "SGLANG_OPT_USE_HIP_MHC_POST": "1", + "SGLANG_OPT_USE_HIP_INT8_SCALED_MM": "0", + "SGLANG_OPT_USE_LMSLIM_INT8_QUANT": "1", + "SGLANG_OPT_USE_W8A8_MARLIN_GEMM": "1", +} + +_PYTHONPATH_EXTRA = "/workspace/sglang/sglang-v0.5.15_k100/python" + + +def _setup_env(): + """Apply K100 env vars and PYTHONPATH once per process.""" + for k, v in _K100_ENV.items(): + if k not in os.environ: + os.environ[k] = v + pp = os.environ.get("PYTHONPATH", "") + if _PYTHONPATH_EXTRA not in pp: + os.environ["PYTHONPATH"] = f"{_PYTHONPATH_EXTRA}:{pp}" if pp else _PYTHONPATH_EXTRA + + +def build_dir_name(args: Dict[str, Any], disable_cuda_graph: bool = False) -> str: + """Build sglang-style directory name from config.""" + parts = [args["version"], f"tp{args['tp_size']}", f"pp{args['pp_size']}"] + parts.append("nograph" if disable_cuda_graph else "graph") + return "_".join(parts) + + +def run_benchmark( + args: Dict[str, Any], + dir_name: str, + batch_size: int, + *, + disable_cuda_graph: bool = False, +) -> bool: + """Run a single bench_one_batch_server invocation.""" + output_dir = os.path.join( + args["output_dir"], "mapping" if disable_cuda_graph else f"bs_{batch_size}" + ) + profile_prefix = f"{dir_name}_" + + cmd = [ + sys.executable, "-m", "sglang.bench_one_batch_server", + "--model-path", args["model_path"], + "--tp-size", str(args["tp_size"]), + "--pp-size", str(args["pp_size"]), + "--batch-size", str(batch_size), + "--input-len", str(args["input_len"]), + "--output-len", str(args["output_len"]), + "--run-name", dir_name, + "--show-report", + "--dataset-name", "random-ids", + "--fake-prefill", + "--profile", + "--profile-start-step", str(args.get("profile_start_step", 5)), + "--profile-steps", str(args.get("profile_steps", 5)), + "--profile-by-stage", + "--profile-prefix", profile_prefix, + "--profile-output-dir", output_dir, + "--disable-radix-cache", + "--chunked-prefill-size", "4096", + "--kv-cache-dtype", "auto", + "--disable-flashinfer-autotune", + "--reasoning-parser", "deepseek-v4", + "--tool-call-parser", "deepseekv4", + "--enable-metrics", + ] + + if disable_cuda_graph: + cmd.append("--disable-cuda-graph") + else: + cmd.extend(["--cuda-graph-bs", str(batch_size)]) + + print(f"\n{'='*80}") + print(f"Running batch_size={batch_size}" + f"{' (CUDA Graph OFF)' if disable_cuda_graph else ''}") + print(f" profile-output-dir: {output_dir}") + print(f" profile-prefix: {profile_prefix}") + print(f"{'='*80}\n") + + try: + subprocess.run(cmd, check=True, timeout=3600) + except subprocess.TimeoutExpired: + print(f"\n[FAILED] batch_size={batch_size}: timed out after 1 hour\n") + return False + except subprocess.CalledProcessError as e: + print(f"\n[FAILED] batch_size={batch_size}: exit code {e.returncode}\n") + return False + return True + + +def main(): + _setup_env() + parser = argparse.ArgumentParser( + description="Run sglang bench_one_batch_server with torch profiler" + ) + parser.add_argument("--config", required=True, + help="Path to JSON benchmark config") + parser.add_argument("--mapping-only", action="store_true", + help="Run only the mapping benchmark (--disable-cuda-graph, one batch)") + parser.add_argument("--formal-only", action="store_true", + help="Run formal benchmarks (CUDA Graph ON, one or all batches)") + parser.add_argument("--single-batch", type=int, default=None, + help="When --formal-only, run only this batch size") + + args = parser.parse_args() + config_path = Path(args.config) + if not config_path.exists(): + print(f"ERROR: config file not found: {args.config}") + return 1 + + with open(config_path) as f: + cfg = json.load(f) + + if args.mapping_only: + dir_name = build_dir_name(cfg, disable_cuda_graph=True) + bs = cfg.get("mapping_batch_size", 8) + ok = run_benchmark(cfg, dir_name, bs, disable_cuda_graph=True) + return 0 if ok else 1 + + if args.formal_only: + dir_name = build_dir_name(cfg, disable_cuda_graph=False) + batch_sizes: List[int] = cfg.get("batch_sizes", [1]) + if args.single_batch is not None: + if args.single_batch in batch_sizes: + batch_sizes = [args.single_batch] + else: + print(f"ERROR: --single-batch {args.single_batch} not in " + f"configured batch_sizes {batch_sizes}") + return 1 + + succeeded, failed = [], [] + for bs in batch_sizes: + ok = run_benchmark(cfg, dir_name, bs) + (succeeded if ok else failed).append(bs) + + print(f"\n{'='*80}") + print(f"Completed: {len(succeeded)} succeeded, {len(failed)} failed") + if succeeded: + print(f" Succeeded batches: {succeeded}") + if failed: + print(f" Failed batches: {failed}") + return 0 if not failed else 1 + + print("ERROR: must specify --mapping-only or --formal-only") + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/metainfer/tasks/sglang_trace_analyze/orchestrator/structure_mapper.py b/metainfer/tasks/sglang_trace_analyze/orchestrator/structure_mapper.py new file mode 100644 index 00000000..b900f9c1 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/orchestrator/structure_mapper.py @@ -0,0 +1,272 @@ +"""Map kernel names (via call stacks) to model structural elements. + +Takes the aggregated kernel list from :mod:`trace_parser` and the model's +``config.json``, then assigns each kernel to: +- ``model_layer`` — e.g. ``layer_{2..58}/attn/qkv_proj`` +- ``op_type`` — GEMM / Attention / Norm / ElementWise / MoE / NCCL / ... +- ``category`` — for grouping (MLA, MoE, GEMM, NCCL, etc.) + +Mapping is done by parsing the Python source location from the call stack +and matching it against known sglang layer source patterns. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional + + +def build_mapping( + kernels: List[Dict[str, Any]], + config: Dict[str, Any], +) -> List[Dict[str, Any]]: + """Build the kernel-to-model-structure mapping. + + Args: + kernels: Aggregated kernel list from ``aggregate_kernels()`` with + ``include_call_stack=True``. + config: Model ``config.json`` as a dict. + + Returns: + List of mapping entries with ``kernel_name``, ``model_layer``, + ``op_type``, ``category``, ``call_stack``, ``confidence``. + """ + entries = [] + for k in kernels: + call_stack = k.get("call_stack", "") + entry = _map_one(k["kernel_name"], call_stack, config) + entries.append(entry) + return entries + + +# ------------------------------------------------------------------ # +# Internal: pattern-based mapping +# ------------------------------------------------------------------ # + +def _map_one( + kernel_name: str, + call_stack: str, + config: Dict[str, Any], + cpu_ops: list | None = None, +) -> Dict[str, Any]: + """Map a single kernel to a model layer by inspecting its call stack + and correlated CPU ops.""" + layer = _infer_layer(call_stack, kernel_name, config, cpu_ops) + op_type = _infer_op_type(kernel_name, call_stack, cpu_ops) + + has_cpu_hint = bool(cpu_ops) + + if not call_stack: + # Confidence tiers without call stack: + # high: kernel name unambiguously identifies op type + # (CK GEMM, flash_attn, fused_moe, NCCL, w8a8, cross_device_reduce) + # medium: CPU ops provide corroborating hint + # low: no useful signal from either source + name_clear = _kernel_name_is_clear(kernel_name, op_type) + cpu_confirms = _cpu_ops_confirm(kernel_name, cpu_ops, op_type) + + if name_clear: + confidence = "high" + elif cpu_confirms: + confidence = "medium" + elif has_cpu_hint: + confidence = "medium" + else: + confidence = "low" + elif layer is None: + confidence = "medium" + + return { + "kernel_name": kernel_name, + "model_layer": layer, + "op_type": op_type, + "category": _op_type_to_category(op_type), + "call_stack": call_stack, + "confidence": confidence, + } + + +def _is_ck_gemm(name: str) -> bool: + """CK (composable_kernel) GEMM kernels have Cijk_ prefix.""" + return name.lower().startswith("cijk_") + + +def _kernel_name_is_clear(kernel_name: str, op_type: str) -> bool: + """Does the kernel name unambiguously identify its op type?""" + name_lower = kernel_name.lower() + # CK GEMM: name encodes tile dims, very clear + if name_lower.startswith("cijk_"): + return True + # Flash attention / MLA kernels + if "flash_fwd" in name_lower or "flash_attn" in name_lower: + return True + # Fused MoE + if "fused_moe" in name_lower: + return True + # NCCL operations + if "nccl" in name_lower: + return True + # w8a8 GEMM kernels (INT8 quantized) + if "w8a8" in name_lower and "scaled_mm" in name_lower: + return True + # Custom allreduce (cross_device_reduce) + if "cross_device_reduce" in name_lower: + return True + # MHC pre/post kernels + if "mhc_pre" in name_lower or "mhc_post" in name_lower: + return True + # topk kernels + if "topk" in name_lower and ("radix" in name_lower or "gather" in name_lower or "find" in name_lower): + return True + return False + + +def _cpu_ops_confirm( + kernel_name: str, + cpu_ops: list | None, + op_type: str, +) -> bool: + """Do the correlated CPU ops confirm the kernel's op type?""" + if not cpu_ops: + return False + cpu_lower = " ".join(cpu_ops).lower() + + confirmations = { + "GEMM": ["aten::linear", "aten::addmm", "aten::matmul", "torch.compile"], + "Attention": ["flash_attn", "flash_fwd", "attention"], + "MoE": ["fused_moe", "moe", "experts"], + "Norm": ["rms_norm", "rmsnorm", "layer_norm", "layernorm"], + "NCCL": ["allreduce", "allgather", "all_reduce", "nccl"], + "Reduce": ["all_reduce", "reduce", "cross_device"], + "ElementWise": ["copy_", "add", "mul", "silu", "gelu", "reshape", "view"], + } + + patterns = confirmations.get(op_type, []) + return any(p in cpu_lower for p in patterns) + + + +def _infer_layer( + call_stack: str, + kernel_name: str, + config: Dict[str, Any], + cpu_ops: list | None = None, +) -> Optional[str]: + """Extract layer information from the call stack and kernel name.""" + name_lower = kernel_name.lower() + cpu_lower = " ".join(cpu_ops or []).lower() + + if call_stack: + import re + lines = call_stack.strip().split("\n") + layer_pat = re.compile(r"model\.layers\.(\d+)") + sglang_layer_pat = re.compile( + r"sglang/srt/layers/(attn|moe|mla|linear|norm|embed|sampler|router)" + ) + for line in lines: + m = layer_pat.search(line) + if m: + return f"layer_{m.group(1)}" + m = sglang_layer_pat.search(line) + if m: + return f"layers/{m.group(1)}" + + # Fallback (no call stack): kernel name + CPU op heuristics + if "flash_fwd" in name_lower or "flash_attn" in name_lower: + return "all_layers/attention" + if "fused_moe" in name_lower or "moe" in cpu_lower: + return "moe_layers/experts" + if name_lower.startswith("cijk_"): + return "all_layers/linear" + if "rms_norm" in cpu_lower or "rmsnorm" in name_lower: + return "all_layers/norm" + if "reduce_kernel" in name_lower: + return "all_layers/allreduce" + if "allgather" in cpu_lower or "nccl" in name_lower: + return "all_layers/communication" + if "elementwise" in name_lower or "vectorized" in name_lower: + return "all_layers/elementwise" + + return None + + +def _infer_op_type(kernel_name: str, call_stack: str, cpu_ops: list | None = None) -> str: + """Infer the op type from kernel name, call stack, and correlated CPU ops. + + Priority: kernel name patterns > CPU op hints > name substring heuristics. + """ + name_lower = kernel_name.lower() + cpu_lower = " ".join(cpu_ops or []).lower() + + # ── Strong kernel name patterns (highest priority) ── + + # CK GEMM kernels (HIP/ROCm composable_kernel) + if name_lower.startswith("cijk_"): + return "GEMM" + + # GPU kernel name patterns — unambiguous from the kernel name itself + if "nccl" in name_lower: + return "NCCL" + if any(k in name_lower for k in ("flash_fwd", "flash_attn")): + return "Attention" + if "fused_moe" in name_lower: + return "MoE" + + # ── Kernel name substring heuristics (medium priority) ── + if "reduce_kernel" in name_lower or "cross_device_reduce" in name_lower: + return "Reduce" + if "elementwise" in name_lower: + return "ElementWise" + if "vectorized" in name_lower: + return "ElementWise" + if "gather" in name_lower: + return "Indexing" + + # ── CPU op hints for torch-compiled/fused kernels ── + if "all_reduce" in cpu_lower: + return "Reduce" # CustomAllReduce, not NCCL + if "allgather" in cpu_lower: + return "NCCL" + if "rms_norm" in cpu_lower: + return "Norm" + + # ── Remaining kernel name patterns (lower priority) ── + if any(k in name_lower for k in ("attn", "attention")): + return "Attention" + if any(k in name_lower for k in ("moe",)): + return "MoE" + if any(k in name_lower for k in ("gemm", "linear", "matmul", "w8a8", "fp8")): + return "GEMM" + if any(k in name_lower for k in ("rmsnorm", "rms_norm", "layernorm")): + return "Norm" + if any(k in name_lower for k in ("allreduce", "allgather", "broadcast")): + return "NCCL" + if any(k in name_lower for k in ("hadamard", "rotate", "rope")): + return "Transform" + if any(k in name_lower for k in ("copy", "memcpy", "memset")): + return "Memory" + if any(k in name_lower for k in ("silu", "gelu", "swiglu", "activation", "act_and_mul")): + return "Activation" + if any(k in name_lower for k in ("topk", "top_k", "gather", "scatter", "sort")): + return "Indexing" + if any(k in name_lower for k in ("quant", "dequant", "fp8_scale")): + return "Quantization" + return "Other" + + +def _op_type_to_category(op_type: str) -> str: + """Map an op_type to a display category.""" + mapping = { + "Attention": "Attention", + "MoE": "MoE", + "GEMM": "GEMM", + "Norm": "Norm", + "NCCL": "NCCL", + "Transform": "Transform", + "Memory": "Memory", + "Activation": "Activation", + "Indexing": "Indexing", + "Quantization": "Quantization", + "Reduce": "Reduce", + "ElementWise": "ElementWise", + } + return mapping.get(op_type, "Other") diff --git a/metainfer/tasks/sglang_trace_analyze/orchestrator/trace_parser.py b/metainfer/tasks/sglang_trace_analyze/orchestrator/trace_parser.py new file mode 100644 index 00000000..7abf9dc0 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/orchestrator/trace_parser.py @@ -0,0 +1,141 @@ +"""Chrome trace JSON parser + kernel aggregation. + +Loads a ``torch.profiler`` Chrome trace (``.json`` or ``.json.gz``) and +produces an aggregated kernel table: one row per unique kernel name, +sorted by total GPU duration descending. + +In the MAPPING phase this also extracts call-stack information for +structure mapping. In the ANALYZE phase it aggregates CUDA Graph replay +events into per-kernel durations. +""" + +from __future__ import annotations + +import gzip +import json +from collections import defaultdict +from pathlib import Path +from typing import Any, Dict, List, Optional + + +def _open_trace(trace_path): + """Open a trace file — transparently handles .gz compression.""" + tp = Path(trace_path) + if tp.suffix == ".gz": + return gzip.open(tp, "rt", encoding="utf-8") + return open(tp, "r", encoding="utf-8") + + +def parse_trace(trace_path) -> Dict[str, Any]: + """Load a Chrome trace JSON and return the top-level document. + + Returns: + Dict with keys: ``traceEvents``, ``displayTimeUnit``, etc. + """ + with _open_trace(trace_path) as f: + data = json.load(f) + return data + + +def aggregate_kernels( + trace_data: Dict[str, Any], + *, + include_call_stack: bool = False, +) -> List[Dict[str, Any]]: + """Aggregate GPU kernel events by kernel name. + + Args: + trace_data: Parsed Chrome trace JSON. + include_call_stack: If True, preserve ``call_stack`` from the first + occurrence of each unique kernel name. + + Returns: + List of kernel dicts sorted by ``total_dur_us`` descending. Each dict: + ``kernel_name``, ``total_dur_us``, ``count``, ``call_stack`` (optional). + """ + trace_events = trace_data.get("traceEvents", []) + if not trace_events: + # sglang sometimes wraps in a list directly + if isinstance(trace_data, list): + trace_events = trace_data + else: + return [] + + # Filter GPU kernel events + kernels: Dict[str, Dict[str, Any]] = {} + for evt in trace_events: + cat = evt.get("cat", "") + name = evt.get("name", "") + dur = evt.get("dur", 0) + + # Torch profiler GPU kernel events: cat="kernel", name like + # "triton_fused_moe_kernel" or "void at::native::..." + if cat != "kernel" or dur <= 0: + continue + + if name not in kernels: + entry: Dict[str, Any] = { + "kernel_name": name, + "total_dur_us": 0, + "count": 0, + } + if include_call_stack: + args = evt.get("args", {}) or {} + call_stack = args.get("call stack", "") + if call_stack: + entry["call_stack"] = call_stack + kernels[name] = entry + + kernels[name]["total_dur_us"] += dur + kernels[name]["count"] += 1 + + # Sort by total duration descending + result = sorted( + kernels.values(), key=lambda k: k["total_dur_us"], reverse=True + ) + return result + + +def aggregate_kernels_with_dims( + trace_data: Dict[str, Any], +) -> List[Dict[str, Any]]: + """Like :func:`aggregate_kernels`, but also collects Input Dims from + ``args["Input Dims"]`` for shape-aware kernels (GEMM, attention). + + This is only meaningful when the trace was captured WITHOUT CUDA Graph + (i.e. during the MAPPING phase), because CUDA Graph replay hides + individual kernel dims. + """ + trace_events = trace_data.get("traceEvents", []) + if isinstance(trace_data, list): + trace_events = trace_data + + kernels: Dict[str, Dict[str, Any]] = {} + for evt in trace_events: + cat = evt.get("cat", "") + name = evt.get("name", "") + dur = evt.get("dur", 0) + if cat != "kernel" or dur <= 0: + continue + + if name not in kernels: + args = evt.get("args", {}) or {} + entry: Dict[str, Any] = { + "kernel_name": name, + "total_dur_us": 0, + "count": 0, + "input_dims": [], + "call_stack": args.get("call stack", ""), + } + kernels[name] = entry + + kernels[name]["total_dur_us"] += dur + kernels[name]["count"] += 1 + args = evt.get("args", {}) or {} + dims = args.get("Input Dims", []) + if dims and dims not in kernels[name]["input_dims"]: + kernels[name]["input_dims"].append(dims) + + return sorted( + kernels.values(), key=lambda k: k["total_dur_us"], reverse=True + ) diff --git a/metainfer/tasks/sglang_trace_analyze/server/__init__.py b/metainfer/tasks/sglang_trace_analyze/server/__init__.py new file mode 100644 index 00000000..3da30f91 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/server/__init__.py @@ -0,0 +1 @@ +"""Server-side plguin for sglang_trace_analyze.""" diff --git a/metainfer/tasks/sglang_trace_analyze/server/_state_readers.py b/metainfer/tasks/sglang_trace_analyze/server/_state_readers.py new file mode 100644 index 00000000..c7e5bb45 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/server/_state_readers.py @@ -0,0 +1,51 @@ +"""State-dir readers for sglang_trace_analyze. + +Reads the authoritative analysis JSON files from +``/analysis/``. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Dict, Optional + + +def _load_json(path: Path) -> Optional[Any]: + if not path.exists(): + return None + try: + return json.loads(path.read_text(encoding="utf-8")) + except (ValueError, OSError): + return None + + +def read_summary(state_dir: Path) -> Optional[Dict[str, Any]]: + return _load_json(state_dir / "analysis" / "summary.json") + + +def read_mapping(state_dir: Path) -> Optional[Dict[str, Any]]: + return _load_json(state_dir / "analysis" / "mapping.json") + + +def read_hints(state_dir: Path) -> Optional[Dict[str, Any]]: + return _load_json(state_dir / "analysis" / "hints.json") + + +def read_batch_detail( + state_dir: Path, bs: int, stage: str +) -> Optional[Dict[str, Any]]: + """Return the combined kernel_table + overlap + fuse for one + (batch_size, stage) pair. + """ + base = state_dir / "analysis" / "batches" / f"bs_{bs}" / stage + kernel_table = _load_json(base / "kernel_table.json") + overlap = _load_json(base / "overlap.json") + fuse = _load_json(base / "fuse.json") + if kernel_table is None and overlap is None and fuse is None: + return None + return { + "kernel_table": kernel_table, + "overlap": overlap, + "fuse": fuse, + } diff --git a/metainfer/tasks/sglang_trace_analyze/server/plugin.py b/metainfer/tasks/sglang_trace_analyze/server/plugin.py new file mode 100644 index 00000000..ea1816f8 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/server/plugin.py @@ -0,0 +1,34 @@ +"""WebPlugin for sglang_trace_analyze — registers routes + detail view.""" + +from __future__ import annotations + +from pathlib import Path + +from metainfer.server.registry import WebPlugin, register + +from .routes import build_router + +PLUGIN_TYPE = "sglang_trace_analyze" +_FRONTEND_DIR = Path(__file__).resolve().parent.parent / "static" +_STATIC_PREFIX = f"/static/plugins/{PLUGIN_TYPE}" + +_IMPORTMAP_ENTRIES: dict = {} + +plugin = WebPlugin( + type=PLUGIN_TYPE, + label="SGLang Trace Analyze", + description=( + "Profile a model with SGLang's torch profiler across multiple batch " + "sizes, then analyze kernel hotspots, TFLOPS/MFU, operator-to-model-" + "structure mapping, fuse opportunities, and generate LLM-powered " + "optimization hints." + ), + build_router=build_router, + detail_view_module="app/sa-detail", + detail_view_export="default", + frontend_dir=_FRONTEND_DIR, + importmap_entries=_IMPORTMAP_ENTRIES, + extra_stylesheets=["sa.css"], +) + +register(plugin) diff --git a/metainfer/tasks/sglang_trace_analyze/server/routes.py b/metainfer/tasks/sglang_trace_analyze/server/routes.py new file mode 100644 index 00000000..3bf73822 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/server/routes.py @@ -0,0 +1,68 @@ +"""FastAPI router for sglang_trace_analyze. + +Routes mounted under ``/api/sglang_trace_analyze/{task_id}``: + + GET /summary → summary.json + GET /mapping → mapping.json + GET /hints → hints.json + GET /batch/{bs}/{stage} → {kernel_table, overlap, fuse} +""" + +from __future__ import annotations + +from fastapi import APIRouter, HTTPException + +from metainfer.server._helpers import ( + require_task_type, + state_dir_for, + task_or_404, +) +from . import _state_readers + +PLUGIN_TYPE = "sglang_trace_analyze" + + +def build_router(plugin) -> APIRouter: + router = APIRouter() + + @router.get("/summary") + def get_summary(task_id: str): + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + data = _state_readers.read_summary(state_dir_for(entry)) + if data is None: + raise HTTPException(404, "summary not yet available") + return data + + @router.get("/mapping") + def get_mapping(task_id: str): + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + data = _state_readers.read_mapping(state_dir_for(entry)) + if data is None: + raise HTTPException(404, "mapping not yet available") + return data + + @router.get("/hints") + def get_hints(task_id: str): + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + data = _state_readers.read_hints(state_dir_for(entry)) + if data is None: + raise HTTPException(404, "hints not yet available") + return data + + @router.get("/batch/{bs}/{stage}") + def get_batch_detail(task_id: str, bs: int, stage: str): + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + data = _state_readers.read_batch_detail( + state_dir_for(entry), bs, stage + ) + if data is None: + raise HTTPException( + 404, f"no analysis data for batch {bs}/{stage}" + ) + return data + + return router diff --git a/metainfer/tasks/sglang_trace_analyze/static/sa-detail.js b/metainfer/tasks/sglang_trace_analyze/static/sa-detail.js new file mode 100644 index 00000000..bf6f8d66 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/static/sa-detail.js @@ -0,0 +1,815 @@ +/** SGLang Trace Analyze — task detail view. + * + * Designed for GPU inference optimization engineers. + * Three tabs: Dashboard | Batch Detail | Optimization Hints + */ +import { html } from "htm/preact"; +import { useEffect, useState, useMemo } from "preact/hooks"; + +const API = (taskId) => `/api/sglang_trace_analyze/${taskId}`; + +export default function SADetail({ taskId }) { + const [summary, setSummary] = useState(null); + const [hints, setHints] = useState(null); + const [detail, setDetail] = useState(null); + const [mapping, setMapping] = useState(null); + const [activeTab, setActiveTab] = useState("dashboard"); + const [activeBatch, setActiveBatch] = useState(null); + const [activeStage, setActiveStage] = useState("decode"); + const [loading, setLoading] = useState(true); + const [error, setError] = useState(null); + + useEffect(() => { + Promise.all([ + fetch(`${API(taskId)}/summary`).then((r) => r.json()), + fetch(`${API(taskId)}/hints`).then((r) => r.json()), + fetch(`${API(taskId)}/mapping`).then((r) => r.json()), + ]) + .then(([s, h, m]) => { setSummary(s); setHints(h); setMapping(m); setLoading(false); }) + .catch((e) => { setError(e.message); setLoading(false); }); + }, [taskId]); + + useEffect(() => { + if (!activeBatch) return; + fetch(`${API(taskId)}/batch/${activeBatch}/${activeStage}`) + .then((r) => r.json()) + .then((d) => setDetail(d)) + .catch(() => setDetail(null)); + }, [taskId, activeBatch, activeStage]); + + if (loading) return html`
Loading analysis…
`; + if (error) return html`
Error: ${error}
`; + if (!summary || !summary.batches || summary.batches.length === 0) { + return html`
No analysis data available yet.
`; + } + + const batchList = summary.batches || []; + if (!activeBatch && batchList.length > 0) setActiveBatch(batchList[0].batch_size); + + return html` +
+
+

Trace Analysis

+ ${summary.model || "?"} | ${summary.gpu || "?"} +
+ +
+ + + +
+ + ${activeTab === "dashboard" && html`<${Dashboard} summary=${summary} detail=${detail} mapping=${mapping} batchList=${batchList} activeBatch=${activeBatch} setActiveBatch=${setActiveBatch} />`} + ${activeTab === "batch" && html` +
+ ${batchList.map((b) => html` + + `)} +
+ ${detail ? html`<${KernelTable} kt=${detail.kernel_table} batch=${activeBatch} stage=${activeStage} />` : html`
Loading…
`} + `} + ${activeTab === "hints" && html`<${HintsPage} hints=${hints} detail=${detail} />`} +
+ `; +} + +/* ═══════════════════════════════════════════════════════════════════════ + DASHBOARD + ═══════════════════════════════════════════════════════════════════════ */ + +function Dashboard({ summary, detail, mapping, batchList, activeBatch, setActiveBatch }) { + if (!detail) return html`
Loading dashboard…
`; + const kt = detail.kernel_table; + if (!kt) return null; + const kernels = kt.kernels || []; + + // Compute stats + const top = kernels[0] || {}; + const cats = {}; + let mfuVals = [], totalDur = kt.total_gpu_time_s || 0; + for (const k of kernels) { + cats[k.category] = (cats[k.category] || 0) + (k.time_pct || 0); + if (k.mfu != null && k.mfu > 0) mfuVals.push(k.mfu); + if (k.tflops_actual != null && k.tflops_actual > 0) mfuVals.push(k.tflops_actual / (k.tflops_theoretical || 192) * 100); + } + const avgMfu = mfuVals.length ? (mfuVals.reduce((a, b) => a + b, 0) / mfuVals.length).toFixed(1) : null; + const computePct = kernels.filter(k => k.bound === "compute").reduce((s, k) => s + (k.time_pct || 0), 0); + const memoryPct = kernels.filter(k => k.bound === "memory").reduce((s, k) => s + (k.time_pct || 0), 0); + const unknownBound = 100 - computePct - memoryPct; + + const ov = detail.overlap || {}; + const cudaGraphOk = (ov.summary || {}).cuda_graph_effective; + const gapCount = (ov.gaps || []).length; + + // Category colors + const catColors = { Reduce: "#c0392b", GEMM: "#d35400", ElementWise: "#e67e22", MoE: "#27ae60", + NCCL: "#e74c3c", Attention: "#8e44ad", Norm: "#2980b9", Indexing: "#16a085", Memory: "#f1c40f", + Quantization: "#2c3e50", Other: "#7f8c8d", Transform: "#2ecc71", Activation: "#e91e63" }; + + const sortedCats = Object.entries(cats).sort((a, b) => b[1] - a[1]); + const totalPct = sortedCats.reduce((s, [, v]) => s + v, 0); + // Build conic-gradient stops for the donut + const donutStops = []; + let acc = 0; + for (const [cat, pct] of sortedCats) { + donutStops.push(`${catColors[cat] || "#95a5a6"} ${acc}% ${acc + pct}%`); + acc += pct; + } + + const summaryText = buildSummary(kernels, cudaGraphOk, top); + + return html` +
+
+ ${summaryText} +
+ + ${/* Row 1: Quick stats */""} +
+
+
${totalDur.toFixed(2)}s
+
Total GPU Time
+
+
+
${(top.time_pct || 0).toFixed(1)}%
+
Top Bottleneck
+
${(top.kernel_name || "").slice(0, 40)}
+
+
+
${avgMfu != null ? avgMfu + "%" : "—"}
+
Avg MFU (BF16)
+
+
+
${cudaGraphOk ? "ON" : "OFF"}
+
CUDA Graph
+
+
+ + ${/* Row 2: Category donut + Bound breakdown + Bottleneck detail */""} +
+
+

GPU Time by Category

+
+
+
+ ${kernels.length} + kernels +
+
+
+ ${sortedCats.slice(0, 8).map(([cat, pct]) => html` +
+ + ${cat} + ${pct.toFixed(1)}% +
+ `)} +
+
+
+ +
+

Bottleneck Detail

+
+
#1
+
+
${top.kernel_name || "?"}
+
+ Category: ${top.category || "?"} | + Op: ${top.op_type || "?"} | + Count: ${top.count || 0} +
+
+ Layer: ${top.model_layer || "unknown"} | + Bound: ${top.bound || "unknown"} | + MFU: ${top.mfu != null ? top.mfu.toFixed(1) + "%" : "—"} +
+
+
+
+
+
+ +

Compute vs Memory Bound

+
+
+ Compute-bound +
+ ${computePct.toFixed(1)}% +
+
+ Memory-bound +
+ ${memoryPct.toFixed(1)}% +
+
+ Unknown +
+ ${unknownBound.toFixed(1)}% +
+
+ +

Overlap

+

${gapCount} idle gaps detected. ${cudaGraphOk ? "CUDA Graph is active — gaps are minimal." : "CUDA Graph is OFF — explore enabling it."}

+
+
+ + ${/* Row 3: TFLOPS & Bandwidth + Structure Mapping */""} +
+ <${TflopsPanel} kernels=${kernels} gpu=${summary.gpu || "K100"} /> + <${StructureMappingPanel} mapping=${mapping} kernels=${kernels} /> +
+ + ${/* Row 4: Fuse + Mapping confidence */""} + <${FusePanel} detail=${detail} /> + + ${/* Row 5: Inefficiency radar + roofline */""} +
+ <${InefficiencyRadar} kernels=${kernels} /> + <${RooflinePanel} kernels=${kernels} gpu=${summary.gpu || "K100"} /> +
+ + ${/* Row 4: Key Findings */""} + <${KeyFindings} kernels=${kernels} kt=${kt} cudaGraph=${cudaGraphOk} /> + + ${/* Row 5: MFU Distribution + Frequency Analysis */""} +
+ <${MfuDistro} kernels=${kernels} gpu=${summary.gpu || "K100"} /> + <${FrequencyPanel} kernels=${kernels} /> +
+ + ${/* Row 6: Top kernels quick preview */""} +
+

Top Kernels

+ + + + ${kernels.slice(0, 10).map((k) => html` + + + + + + + + + + + `)} + +
#%CategoryKernelCountAvg μsMFUBound
${k.rank}
${(k.time_pct || 0).toFixed(1)}%
${k.category || "?"}${(k.kernel_name || "").slice(0, 55)}${k.count}${(k.avg_dur_us || 0).toFixed(1)}${k.mfu != null ? k.mfu.toFixed(1) + "%" : "—"}${k.bound || "?"}
+
+
+ `; +} + +/* ═══════════════════════════════════════════════════════════════════════ + KERNEL TABLE (full, searchable) + ═══════════════════════════════════════════════════════════════════════ */ + +function KernelTable({ kt, batch, stage }) { + if (!kt) return null; + const kernels = kt.kernels || []; + const totalTime = kt.total_gpu_time_s || 1; + const [search, setSearch] = useState(""); + const [catFilter, setCatFilter] = useState("all"); + + const categories = [...new Set(kernels.map((k) => k.category || "Other"))]; + const filtered = kernels.filter((k) => { + if (catFilter !== "all" && k.category !== catFilter) return false; + if (search && !k.kernel_name.toLowerCase().includes(search.toLowerCase())) return false; + return true; + }); + + return html` +
+

Kernel Hotspots — BS=${batch} ${stage} (${kernels.length} unique, ${totalTime.toFixed(1)}s GPU)

+
+ setSearch(e.target.value)} /> + + ${filtered.length} of ${kernels.length} kernels +
+
+ + + + ${filtered.slice(0, 100).map((k) => html` + + + + + + + + + + + + + + `)} + +
#%CategoryOpLayerCountAvg μsMFUBoundConfKernel
${k.rank}
${(k.time_pct || 0).toFixed(1)}%
${k.category || "?"}${k.op_type || "?"}${k.model_layer || "—"}${k.count}${(k.avg_dur_us || 0).toFixed(1)}${k.mfu != null ? k.mfu.toFixed(1) + "%" : "—"}${k.bound || "—"}${k.confidence || "?"}${(k.kernel_name || "").slice(0, 60)}
+
+
+ `; +} + +/* ═══════════════════════════════════════════════════════════════════════ + HINTS PAGE + ═══════════════════════════════════════════════════════════════════════ */ + +function HintsPage({ hints, detail }) { + const kt = detail ? detail.kernel_table : null; + const fuse = detail ? detail.fuse : null; + const fuseMatches = (fuse && fuse.matches) || []; + + return html` + ${kt && html` + <${BottleneckAnalysis} kt=${kt} /> + `} + + ${fuseMatches.length > 0 && html` +
+

Fuse Opportunities (${fuseMatches.length})

+ ${fuseMatches.map((m) => html` +
+ ${m.pattern} + ${m.confidence} + ~${m.estimated_saving_us}μs saving +

${m.suggestion}

+
+ `)} +
+ `} + + ${hints && hints.status !== "skipped" && html` +
+

AI Optimization Hints

+ ${(hints.suggestions || []).map((s) => html` +
+ ${s.title} + ${s.difficulty} + Est. saving: ${s.estimated_saving_pct}% +

${s.what_to_change}

+

${s.why} | Type: ${s.category}

+
+ `)} +
+ `} + + ${(hints && hints.status === "skipped" && fuseMatches.length === 0 && !kt) && html` +

Optimization Hints

No hints or fuse matches available yet.

+ `} + `; +} + +function BottleneckAnalysis({ kt }) { + if (!kt) return null; + const kernels = kt.kernels || []; + const top = kernels[0]; + const top3 = kernels.slice(0, 3); + + const computeBoundPct = kernels.filter(k => k.bound === "compute").reduce((s, k) => s + (k.time_pct || 0), 0); + const suggestions = []; + if (computeBoundPct < 30) suggestions.push("Most kernels are memory-bound — focus on kernel fusion to reduce memory traffic."); + if ((top.time_pct || 0) > 50) suggestions.push(`"${(top.kernel_name || "").slice(0, 40)}" dominates at ${(top.time_pct || 0).toFixed(1)}%. Consider optimizing or replacing this kernel.`); + if (suggestions.length === 0) suggestions.push("GPU time is spread across many kernels. Look for fusion opportunities in the table below."); + + return html` +
+

Bottleneck Analysis

+
+ ${top3.map((k, i) => html` +
+ #${i + 1} + ${(k.time_pct || 0).toFixed(1)}% + ${(k.kernel_name || "").slice(0, 60)} + ${k.category || "?"} +
+ `)} +
+ ${suggestions.map((s) => html`

${s}

`)} +
+ `; +} + +/* ── Inefficiency Radar: high-time, low-MFU kernels ── */ + +function InefficiencyRadar({ kernels }) { + if (!kernels || !kernels.length) return null; + // Top kernels by (time_pct * (100 - mfu)) / 100 — high time, low efficiency + const inefficiency = kernels + .filter((k) => (k.time_pct || 0) > 0.03) + .map((k) => ({ + ...k, + waste: ((k.time_pct || 0) * (k.mfu != null ? Math.max(0, 100 - k.mfu) : 100)) / 100, + })) + .sort((a, b) => b.waste - a.waste); + + return html` +
+

Inefficiency Radar

+

Kernels with high GPU time and low MFU — biggest optimization potential.

+ + + + ${inefficiency.slice(0, 8).map((k) => html` + + + + + + + + `)} + +
KernelTime%MFUWaste ScoreCategory
${(k.kernel_name || "").slice(0, 50)}${(k.time_pct || 0).toFixed(1)}%${k.mfu != null ? k.mfu.toFixed(1) + "%" : "—"}${k.waste.toFixed(1).replace(/^-/, "")}${k.category || "?"}
+
+ `; +} + +/* ── Roofline Analysis ── */ + +function RooflinePanel({ kernels, gpu }) { + if (!kernels || !kernels.length) return null; + + // GPU peaks + const peaks = { K100: { bf16: 192, bw: 700 }, A100_80G: { bf16: 312, bw: 2039 }, + H100: { bf16: 989, bw: 3350 }, B200: { bf16: 2250, bw: 8000 } }; + const pk = peaks[gpu] || peaks.K100; + const peakFlops = pk.bf16 * 1e12; // TFLOPS → FLOPS + const peakBw = pk.bw * 1e9; // GB/s → B/s + const ridgePoint = peakFlops / peakBw; // ops/byte at the ridge + + // Classify each kernel with valid data + const pts = kernels + .filter((k) => k.tflops_actual != null && k.tflops_actual > 0 && k.bandwidth_gb_s != null && k.bandwidth_gb_s > 0) + .map((k) => ({ + name: k.kernel_name, category: k.category, time_pct: k.time_pct, + flops: k.tflops_actual * 1e12, bw: k.bandwidth_gb_s * 1e9, + opsPerByte: (k.tflops_actual * 1e12) / (k.bandwidth_gb_s * 1e9), + bound: k.bound, rank: k.rank, + })); + + const computeBound = pts.filter((p) => p.bound === "compute").length; + const memoryBound = pts.filter((p) => p.bound === "memory").length; + + return html` +
+

Roofline Analysis

+

GPU: ${gpu} | Peak BF16: ${pk.bf16} TFLOPS | BW: ${pk.bw} GB/s | Ridge: ${ridgePoint.toFixed(0)} ops/byte

+

+ ${computeBound} compute-bound | + ${memoryBound} memory-bound + ${pts.length < 5 ? html` (${kernels.length - pts.length} kernels lack dims for roofline)` : ""} +

+
+ ${pts.slice(0, 12).map((p) => { + const barW = Math.min(Math.log10(Math.max(p.opsPerByte, 1)) / Math.log10(ridgePoint * 10) * 100, 100); + const onRidge = p.opsPerByte > ridgePoint; + return html` +
+ ${(p.name || "").slice(0, 40)} + +
+ + ${p.opsPerByte.toFixed(0)} op/B + ${onRidge ? "compute" : "memory"} +
+ `; + })} +
+

Ridge point: ${ridgePoint.toFixed(0)} ops/byte. Left of ridge = memory-bound. Right = compute-bound.

+
+ `; +} + +/* ── Executive Summary builder ── */ + +function buildSummary(kernels, cudaGraph, top) { + if (!kernels || !kernels.length) return "No analysis data available."; + + const parts = []; + parts.push(cudaGraph ? "CUDA Graph ON" : "CUDA Graph OFF"); + + if (top && top.category) { + parts.push(`${top.category} is your bottleneck (${(top.time_pct || 0).toFixed(0)}%)`); + } + + // Find category insights + const cats = {}; + for (const k of kernels) cats[k.category] = (cats[k.category] || 0) + (k.time_pct || 0); + + const reducePct = cats["Reduce"] || 0; + if (reducePct < 10 && reducePct > 0) { + parts.push(`TP allreduce well-optimized (${reducePct.toFixed(0)}%)`); + } + + const gemmPct = cats["GEMM"] || 0; + if (gemmPct > 20) { + parts.push(`quantize GEMMs to FP8 for ~${(gemmPct * 0.4).toFixed(0)}% improvement`); + } + + const elementPct = cats["ElementWise"] || 0; + if (elementPct > 10) { + parts.push(`fuse element-wise ops to save ~${(elementPct * 0.3).toFixed(0)}%`); + } + + if (!cudaGraph) { + parts.push("enable CUDA Graph for 3-5x speedup"); + } + + return parts.join(". ") + "."; +} + +/* ── Key Findings auto-summary ── */ + +function KeyFindings({ kernels, kt, cudaGraph }) { + if (!kernels || !kernels.length) return null; + + const total = kt.total_gpu_time_s || 0; + const top = kernels[0]; + const top3 = kernels.slice(0, 3); + + // Build findings from data + const findings = []; + + // 1. Dominant kernel + if ((top.time_pct || 0) > 30) { + findings.push({ + icon: "🔴", title: "Single kernel dominates", + text: `"${(top.kernel_name || "").slice(0, 45)}" consumes ${(top.time_pct || 0).toFixed(1)}% of GPU time alone. This is your primary optimization target.`, + }); + } else if ((top.time_pct || 0) > 15) { + findings.push({ + icon: "🟡", title: "Moderate hotspot", + text: `Top kernel "${(top.kernel_name || "").slice(0, 45)}" at ${(top.time_pct || 0).toFixed(1)}%. Consider fusion or replacement.`, + }); + } else { + findings.push({ + icon: "🟢", title: "Well-distributed workload", + text: "GPU time is spread across many kernels. Focus on fusion and reducing kernel launch overhead.", + }); + } + + // 2. CUDA Graph + if (cudaGraph) { + findings.push({ + icon: "🟢", title: "CUDA Graph active", + text: `Total GPU time: ${total.toFixed(2)}s with CUDA Graph. Kernel launch overhead is minimized.`, + }); + } else { + findings.push({ + icon: "🔴", title: "CUDA Graph disabled", + text: "Enable CUDA Graph to reduce kernel launch overhead and CPU-GPU synchronization. Expected 3-5x speedup on decode.", + }); + } + + // 3. Category concentration + const cats = {}; + for (const k of kernels) cats[k.category] = (cats[k.category] || 0) + (k.time_pct || 0); + const topCat = Object.entries(cats).sort((a, b) => b[1] - a[1])[0]; + if (topCat && topCat[1] > 50) { + findings.push({ + icon: "🔴", title: `Category "${topCat[0]}" dominates at ${topCat[1].toFixed(0)}%`, + text: topCat[0] === "Reduce" ? "TP allreduce is the bottleneck. Consider communication-computation overlap or reducing TP degree." : + topCat[0] === "GEMM" ? "GEMM is the bottleneck. Explore quantization (FP8/INT8) or faster GEMM backends." : + `Focus optimization efforts on ${topCat[0]} operations.`, + }); + } + + // 4. Top 3 summary + const top3Summary = top3.map((k, i) => + `#${i + 1} ${(k.category || "?").slice(0, 10)} ${(k.time_pct || 0).toFixed(1)}%` + ).join(" | "); + findings.push({ + icon: "📊", title: "Top 3 kernels", + text: top3Summary, + }); + + // 5. MFU note + const withMfu = kernels.filter((k) => k.mfu != null && k.mfu > 0); + if (withMfu.length === 0) { + findings.push({ + icon: "💡", title: "No MFU data available", + text: "Profiler was run without record_shapes=True. Enable it to get per-kernel TFLOPS and MFU analysis.", + }); + } else if (withMfu.length < 10) { + findings.push({ + icon: "💡", title: `MFU data available for ${withMfu.length} kernels`, + text: "Limited TFLOPS data (only CK GEMM tiles). Enable record_shapes=True for full MFU coverage.", + }); + } + + return html` +
+

Key Findings

+
+ ${findings.map((f) => html` +
+ ${f.icon} +
+ ${f.title} +

${f.text}

+
+
+ `)} +
+
+ `; +} + +/* ── MFU Distribution Histogram ── */ + +function MfuDistro({ kernels, gpu }) { + if (!kernels || !kernels.length) return null; + const peaks = { K100: { bf16: 192 }, A100_80G: { bf16: 312 }, H100: { bf16: 989 }, B200: { bf16: 2250 } }; + const pk = (peaks[gpu] || peaks.K100).bf16; + + // Compute MFU for ALL kernels from tflops_actual / theoretical + const mfuVals = kernels.map((k) => { + if (k.mfu != null) return k.mfu; + if (k.tflops_actual != null && k.tflops_actual > 0) return k.tflops_actual / pk * 100; + return null; + }).filter((v) => v != null); + + if (mfuVals.length === 0) return html`

MFU Distribution

No MFU data available (no Input Dims in trace).

`; + + const buckets = [0, 5, 10, 25, 50, 75, 90, 100]; + const labels = ["0-5%", "5-10%", "10-25%", "25-50%", "50-75%", "75-90%", "90-100%"]; + const hist = new Array(buckets.length - 1).fill(0); + for (const v of mfuVals) { + for (let i = buckets.length - 1; i >= 0; i--) { + if (v >= buckets[i]) { hist[i]++; break; } + } + } + + const maxN = Math.max(...hist, 1); + const avg = mfuVals.reduce((a, b) => a + b, 0) / mfuVals.length; + const median = mfuVals.sort((a, b) => a - b)[Math.floor(mfuVals.length / 2)]; + + return html` +
+

MFU Distribution

+

${mfuVals.length} kernels with TFLOPS data | avg=${avg.toFixed(1)}% | median=${median.toFixed(1)}%

+
+ ${hist.map((n, i) => html` +
+ ${labels[i]} +
+
+
+ ${n} +
+ `)} +
+
+ `; +} + +/* ── Frequency Analysis ── */ + +function FrequencyPanel({ kernels }) { + if (!kernels || !kernels.length) return null; + // Top kernels by call count + const byCount = [...kernels].sort((a, b) => (b.count || 0) - (a.count || 0)); + + return html` +
+

Top by Invocation Count

+

High invocation count kernels may indicate repeated small operations that could be batched.

+ + + + ${byCount.slice(0, 10).map((k) => html` + + + + + + + + `)} + +
KernelCallsTime%Avg μsCategory
${(k.kernel_name || "").slice(0, 45)}${k.count}${(k.time_pct || 0).toFixed(1)}%${(k.avg_dur_us || 0).toFixed(1)}${k.category || "?"}
+
+ `; +} + +/* ── TFLOPS & Bandwidth Panel ── */ + +function TflopsPanel({ kernels, gpu }) { + if (!kernels || !kernels.length) return null; + const peaks = { K100: { bf16: 192, bw: 700 }, A100_80G: { bf16: 312, bw: 2039 }, + H100: { bf16: 989, bw: 3350 }, B200: { bf16: 2250, bw: 8000 } }; + const pk = peaks[gpu] || peaks.K100; + + // Kernels with actual TFLOPS data + const withData = kernels.filter((k) => k.tflops_actual != null && k.tflops_actual > 0); + const withBw = kernels.filter((k) => k.bandwidth_gb_s != null && k.bandwidth_gb_s > 0); + + return html` +
+

TFLOPS & Bandwidth

+

GPU: ${gpu} | Theoretical peak BF16: ${pk.bf16} TFLOPS | BW: ${pk.bw} GB/s

+

${withData.length}/${kernels.length} kernels have TFLOPS data (CK GEMM tile dims extracted from kernel names).

+ + + + ${kernels.filter(k => k.tflops_actual != null || k.bandwidth_gb_s != null).slice(0, 10).map((k) => html` + + + + + + + + + `)} + +
KernelTFLOPSPeak%BW GB/sBW%Bound
${(k.kernel_name || "").slice(0, 45)}${k.tflops_actual != null ? k.tflops_actual.toFixed(3) : "—"}${k.mfu != null ? k.mfu.toFixed(1) + "%" : "—"}${k.bandwidth_gb_s != null ? k.bandwidth_gb_s.toFixed(1) : "—"}${k.bandwidth_gb_s != null ? (k.bandwidth_gb_s / pk.bw * 100).toFixed(1) + "%" : "—"}${k.bound || "—"}
+
+ `; +} + +/* ── Model Structure → Operator Mapping Panel ── */ + +function StructureMappingPanel({ mapping, kernels }) { + if (!mapping || !mapping.entries) return html`

Model Structure Mapping

No mapping data available.

`; + + const entries = mapping.entries || []; + // Group by model_layer + const layerGroups = {}; + for (const e of entries) { + const layer = e.model_layer || "unknown"; + if (!layerGroups[layer]) layerGroups[layer] = { kernels: [], categories: {} }; + layerGroups[layer].kernels.push(e); + layerGroups[layer].categories[e.category] = (layerGroups[layer].categories[e.category] || 0) + 1; + } + + const layers = Object.entries(layerGroups).sort((a, b) => b[1].kernels.length - a[1].kernels.length); + + // Confidence stats + const confStats = { high: 0, medium: 0, low: 0 }; + for (const e of entries) { confStats[e.confidence || "low"]++; } + const total = entries.length || 1; + + return html` +
+

Model Structure → Operator Mapping

+

${entries.length} kernel↔layer mappings | + high ${confStats.high} (${(confStats.high/total*100).toFixed(0)}%) + med ${confStats.medium} (${(confStats.medium/total*100).toFixed(0)}%) + low ${confStats.low} (${(confStats.low/total*100).toFixed(0)}%) +

+
+ ${layers.slice(0, 10).map(([layer, group]) => html` +
+ ${layer} + ${group.kernels.length} kernels + + ${Object.entries(group.categories).slice(0, 4).map(([cat, n]) => html` + ${cat}×${n} + `)} + +
+ `)} +
+
+ `; +} + +/* ── Fuse Opportunities Panel ── */ + +function FusePanel({ detail }) { + const fuse = detail ? detail.fuse : null; + const matches = fuse ? (fuse.matches || []) : []; + + if (matches.length === 0) return html` +
+

Fuse Opportunities

+

No fuse pattern matches found in rule engine. Try enabling LLM hints for AI-generated suggestions.

+
+ `; + + return html` +
+

Fuse Opportunities (${matches.length})

+ ${matches.map((m) => html` +
+
+ ${m.pattern} + ${m.confidence} + ~${m.estimated_saving_us}μs estimated saving +
+

${m.suggestion}

+

Kernels: ${(m.kernels || []).join(" → ")}

+
+ `)} +
+ `; +} diff --git a/metainfer/tasks/sglang_trace_analyze/static/sa.css b/metainfer/tasks/sglang_trace_analyze/static/sa.css new file mode 100644 index 00000000..097e512b --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/static/sa.css @@ -0,0 +1,172 @@ +/* sglang_trace_analyze — dashboard + kernel table styles */ +.sa-detail { padding: 12px 16px; color: #e0e0e0; font-family: system-ui, sans-serif; } +.sa-header { margin-bottom: 8px; } +.sa-header h2 { margin: 0 0 2px; color: #fff; font-size: 18px; } +.sa-meta { color: #888; font-size: 12px; } +.sa-loading,.sa-error,.sa-empty { padding: 32px; text-align: center; color: #888; } +.sa-error { color: #e74c3c; } + +/* Tabs */ +.sa-tabs { display: flex; gap: 2px; margin-bottom: 10px; border-bottom: 2px solid #333; } +.sa-tab-btn { padding: 6px 16px; border: none; border-radius: 4px 4px 0 0; + background: transparent; color: #999; cursor: pointer; font-size: 13px; } +.sa-tab-btn.active-tab { background: #2a2a2a; color: #4a90d9; font-weight: 600; } +.sa-batch-tabs { display: flex; gap: 6px; margin-bottom: 10px; } +.sa-tab { padding: 4px 12px; border: 1px solid #444; border-radius: 4px; + background: #2a2a2a; color: #ccc; cursor: pointer; font-size: 12px; } +.sa-tab.active { background: #4a90d9; color: #fff; border-color: #4a90d9; } + +/* Panels */ +.sa-panel { background: #1e1e1e; border: 1px solid #333; border-radius: 6px; + padding: 14px; margin-bottom: 10px; } +.sa-panel h3 { margin: 0 0 8px; color: #ddd; font-size: 13px; } +.sa-note { color: #888; font-size: 11px; margin: 4px 0; } +.ml8 { margin-left: 8px; } + +/* Summary banner */ +.sa-summary-banner { background: linear-gradient(135deg, #1a2a3a 0%, #1e1e1e 100%); + border: 1px solid #4a90d9; border-radius: 6px; padding: 10px 14px; + margin-bottom: 10px; font-size: 13px; color: #ddd; line-height: 1.5; } + +/* Stat cards */ +.sa-stats-row { display: grid; grid-template-columns: repeat(4, 1fr); gap: 10px; margin-bottom: 10px; } +.sa-stat-card { background: #1e1e1e; border: 1px solid #333; border-radius: 6px; + padding: 12px; text-align: center; } +.sa-stat-value { font-size: 24px; font-weight: 700; color: #fff; line-height: 1.2; } +.sa-stat-label { font-size: 11px; color: #888; margin-top: 4px; } +.sa-stat-sub { font-size: 10px; color: #666; margin-top: 2px; + overflow: hidden; text-overflow: ellipsis; white-space: nowrap; max-width: 180px; } +.sa-stat-warn { border-color: #c0392b; } +.sa-stat-warn .sa-stat-value { color: #e74c3c; } +.sa-stat-ok { border-color: #27ae60; } +.sa-stat-ok .sa-stat-value { color: #2ecc71; } + +/* Grid */ +.sa-grid-2col { display: grid; grid-template-columns: 380px 1fr; gap: 10px; margin-bottom: 10px; } + +/* Donut chart */ +.sa-donut-wrap { display: flex; align-items: center; gap: 16px; } +.sa-donut { width: 140px; height: 140px; border-radius: 50%; position: relative; flex-shrink: 0; } +.sa-donut-hole { position: absolute; top: 28px; left: 28px; right: 28px; bottom: 28px; + background: #1e1e1e; border-radius: 50%; display: flex; flex-direction: column; + align-items: center; justify-content: center; } +.sa-donut-val { font-size: 22px; font-weight: 700; color: #fff; } +.sa-donut-lbl { font-size: 10px; color: #888; } +.sa-donut-legend { flex: 1; min-width: 0; } +.sa-legend-item { display: flex; align-items: center; gap: 6px; margin-bottom: 3px; } +.sa-legend-swatch { width: 10px; height: 10px; border-radius: 2px; flex-shrink: 0; } +.sa-legend-name { font-size: 11px; color: #ccc; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } +.sa-legend-pct { font-size: 11px; color: #aaa; margin-left: auto; font-variant-numeric: tabular-nums; } + +/* Bottleneck */ +.sa-bottleneck { display: flex; gap: 12px; } +.sa-bn-rank { font-size: 32px; font-weight: 800; color: #c0392b; line-height: 1; flex-shrink: 0; } +.sa-bn-info { flex: 1; min-width: 0; } +.sa-bn-name { font-size: 12px; font-family: monospace; color: #e74c3c; margin-bottom: 4px; + word-break: break-all; } +.sa-bn-meta { font-size: 11px; color: #888; margin-bottom: 2px; } +.sa-bn-meta strong { color: #ccc; } +.sa-bn-bar-wrap { background: #2a2a2a; border-radius: 3px; height: 20px; overflow: hidden; margin-top: 6px; } +.sa-bn-bar { height: 100%; background: #c0392b; border-radius: 3px; min-width: 2px; } + +/* Bound bars */ +.sa-bound-bars { } +.sa-bound-row { display: flex; align-items: center; gap: 8px; margin-bottom: 5px; } +.sa-bound-label { width: 100px; font-size: 11px; color: #ccc; text-align: right; flex-shrink: 0; } +.sa-bound-bar-bg { flex: 1; background: #2a2a2a; border-radius: 3px; height: 14px; overflow: hidden; } +.sa-bound-bar { height: 100%; border-radius: 3px; min-width: 2px; } +.sa-bb-compute { background: #2ecc71; } +.sa-bb-memory { background: #e67e22; } +.sa-bb-unknown { background: #555; } +.sa-bound-pct { width: 45px; font-size: 11px; color: #aaa; font-variant-numeric: tabular-nums; } + +/* Bottleneck list (Hints page) */ +.sa-bn-list { margin-bottom: 10px; } +.sa-bn-row { display: flex; align-items: center; gap: 8px; padding: 4px 0; border-bottom: 1px solid #2a2a2a; } +.sa-bn-rank-sm { font-size: 12px; font-weight: 700; color: #888; min-width: 24px; } +.sa-bn-pct { font-size: 13px; font-weight: 600; color: #e74c3c; min-width: 48px; font-variant-numeric: tabular-nums; } +.sa-bn-name-sm { font-size: 11px; font-family: monospace; color: #ccc; flex: 1; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } + +/* Filter bar */ +.sa-filters { display: flex; gap: 8px; align-items: center; margin-bottom: 8px; } +.sa-search { background: #2a2a2a; border: 1px solid #444; border-radius: 4px; padding: 4px 8px; + color: #ddd; font-size: 12px; width: 220px; } +.sa-select { background: #2a2a2a; border: 1px solid #444; border-radius: 4px; padding: 4px 8px; + color: #ddd; font-size: 12px; } +.sa-filter-count { font-size: 11px; color: #888; } + +/* Table */ +.sa-table-wrap { overflow-x: auto; } +.sa-table { width: 100%; border-collapse: collapse; font-size: 11px; } +.sa-table th { text-align: left; padding: 5px 6px; border-bottom: 1px solid #333; + color: #999; font-weight: 600; white-space: nowrap; } +.sa-table td { padding: 3px 6px; border-bottom: 1px solid #2a2a2a; vertical-align: middle; } +.sa-num { text-align: right; font-variant-numeric: tabular-nums; color: #aaa; } +.sa-sm { font-size: 10px; color: #888; } +.sa-kernel-name { max-width: 300px; overflow: hidden; text-overflow: ellipsis; + white-space: nowrap; font-family: monospace; font-size: 10px; color: #bbb; } +.sa-pct { width: 100px; } +.sa-bar-bg { position: relative; background: #2a2a2a; border-radius: 2px; + height: 14px; overflow: hidden; } +.sa-bar { position: absolute; left: 0; top: 0; height: 100%; + background: #4a90d9; border-radius: 2px; opacity: 0.5; } +.sa-bar-bg span { position: relative; z-index: 1; font-size: 10px; + line-height: 14px; padding-left: 3px; color: #ddd; } +.sa-cat { display: inline-block; padding: 1px 5px; border-radius: 2px; + font-size: 10px; background: #333; color: #ccc; } + +/* Confidence badges */ +.sa-conf { display: inline-block; padding: 1px 4px; border-radius: 2px; font-size: 10px; } +.sa-conf-high { background: #27ae60; color: #fff; } +.sa-conf-medium { background: #e67e22; color: #fff; } +.sa-conf-low { background: #c0392b; color: #fff; } + +/* Hints */ +.sa-hint-card { background: #252525; border-left: 3px solid #4a90d9; + padding: 8px 10px; margin-bottom: 8px; border-radius: 0 4px 4px 0; } +.sa-hint-card strong { color: #f1c40f; } +.sa-difficulty { display: inline-block; padding: 1px 6px; border-radius: 3px; + font-size: 10px; color: #fff; margin-left: 6px; } +.sa-diff-low { background: #27ae60; } +.sa-diff-medium { background: #e67e22; } +.sa-diff-high { background: #c0392b; } +.sa-suggestion { font-size: 12px; color: #f1c40f; background: #2a2a20; + border-left: 3px solid #f1c40f; padding: 6px 10px; margin: 6px 0; border-radius: 0 4px 4px 0; } + +/* Roofline */ +.sa-roofline-bars { margin-top: 8px; } +.sa-rf-row { display: flex; align-items: center; gap: 6px; margin-bottom: 4px; } +.sa-rf-name { font-size: 10px; font-family: monospace; color: #bbb; width: 160px; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; flex-shrink: 0; } +.sa-rf-bar-wrap { flex: 1; background: #2a2a2a; border-radius: 2px; height: 12px; overflow: hidden; } +.sa-rf-bar { height: 100%; border-radius: 2px; min-width: 2px; } +.sa-rf-compute { background: #2ecc71; } +.sa-rf-memory { background: #e67e22; } +.sa-rf-val { font-size: 10px; color: #888; width: 65px; text-align: right; font-variant-numeric: tabular-nums; flex-shrink: 0; } +.sa-rf-bound { font-size: 10px; color: #666; width: 60px; flex-shrink: 0; } + +/* Mapping grid */ +.sa-mapping-grid { margin-top: 8px; } +.sa-mapping-row { display: flex; align-items: center; gap: 8px; padding: 3px 0; border-bottom: 1px solid #2a2a2a; } +.sa-mapping-layer { font-size: 11px; color: #ccc; min-width: 140px; font-family: monospace; } +.sa-mapping-count { font-size: 10px; color: #888; min-width: 60px; } +.sa-mapping-cats { display: flex; gap: 4px; flex-wrap: wrap; } + +/* Fuse cards */ +.sa-fuse-card { background: #252525; border-left: 3px solid #e67e22; padding: 8px 10px; margin-bottom: 6px; border-radius: 0 4px 4px 0; } +.sa-fuse-header { display: flex; align-items: center; gap: 8px; margin-bottom: 4px; } +.sa-fuse-header strong { color: #f1c40f; } + +/* MFU Histogram */ +.sa-hist { margin-top: 6px; } +.sa-hist-row { display: flex; align-items: center; gap: 8px; margin-bottom: 4px; } +.sa-hist-label { font-size: 10px; color: #888; width: 55px; text-align: right; flex-shrink: 0; } +.sa-hist-bar-bg { flex: 1; background: #2a2a2a; border-radius: 2px; height: 14px; overflow: hidden; } +.sa-hist-bar { height: 100%; background: #4a90d9; border-radius: 2px; min-width: 2px; } +.sa-hist-count { font-size: 10px; color: #aaa; width: 30px; text-align: right; font-variant-numeric: tabular-nums; flex-shrink: 0; } + +/* Key Findings */ +.sa-findings { display: flex; flex-wrap: wrap; gap: 8px; } +.sa-finding-card { display: flex; gap: 8px; background: #252525; border-radius: 4px; padding: 8px 10px; flex: 1; min-width: 280px; max-width: calc(50% - 4px); } +.sa-finding-icon { font-size: 16px; flex-shrink: 0; line-height: 1.2; } +.sa-finding-body { min-width: 0; } +.sa-finding-body strong { font-size: 12px; color: #ddd; } diff --git a/metainfer/tasks/sglang_trace_analyze/tests/__init__.py b/metainfer/tasks/sglang_trace_analyze/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/metainfer/tasks/sglang_trace_analyze/tests/test_flops_calculator.py b/metainfer/tasks/sglang_trace_analyze/tests/test_flops_calculator.py new file mode 100644 index 00000000..ef31820b --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/tests/test_flops_calculator.py @@ -0,0 +1,78 @@ +"""FLOPs calculator tests.""" + +from ..orchestrator.gpu_specs import GpuSpec +from ..orchestrator.flops_calculator import ( + _estimate_flops, + _estimate_bytes, + calculate_mfu, +) + + +K100 = GpuSpec( + label="K100", + fp32_tflops=49, + tf32_tflops=98, + bf16_tflops=192, + fp16_tflops=192, + int8_tops=392, + bandwidth_gb_s=700, +) + + +def test_estimate_flops_gemm_3d(): + # M=4096, K=2048, N=512 → 2*4096*2048*512 = 8,589,934,592 + flops = _estimate_flops("GEMM", [[4096, 2048, 512]], batch_size=1) + assert flops == 2 * 4096 * 2048 * 512 + + +def test_estimate_flops_gemm_batched(): + # B=4, M=1024, N=512, K=2048 → 2*4*1024*2048*512 + flops = _estimate_flops("GEMM", [[4, 1024, 512, 2048]], batch_size=4) + assert flops == 2 * 4 * 1024 * 2048 * 512 + + +def test_estimate_flops_no_dims(): + assert _estimate_flops("GEMM", [], batch_size=8) == 0 + + +def test_estimate_bytes_gemm(): + bytes_moved = _estimate_bytes("GEMM", [[4096, 2048, 512]], batch_size=1) + # (4096*2048 + 2048*512 + 4096*512) * 2 bytes + expected = (4096 * 2048 + 2048 * 512 + 4096 * 512) * 2 + assert bytes_moved == expected + + +def test_calculate_mfu_basic(): + # 2*4096*2048*512 = 8.59e9 FLOPs. At 10 us this is ~859 TFLOPS + # (far above K100 peak), but this is synthetic — we just verify + # the fields are populated and reasonable. + kernels = [ + { + "kernel_name": "triton_gemm", + "total_dur_us": 50, # 50 us for 8.6e9 FLOPs = 172 TFLOPS + "count": 1, + "input_dims": [[4096, 2048, 512]], + "op_type": "GEMM", + } + ] + result = calculate_mfu(kernels, K100, batch_size=1, dtype="bf16") + k = result[0] + assert k["tflops_theoretical"] == 192 + assert k["tflops_actual"] > 0 + assert k["mfu"] > 0 + assert k["bound"] in ("compute", "memory") + + +def test_calculate_mfu_no_dims(): + kernels = [ + { + "kernel_name": "cuda_graph_replay", + "total_dur_us": 500_000, + "count": 1, + "input_dims": [], + "op_type": "Other", + } + ] + result = calculate_mfu(kernels, K100, batch_size=8, dtype="bf16") + assert result[0]["tflops_actual"] == 0 + assert result[0]["mfu"] == 0 diff --git a/metainfer/tasks/sglang_trace_analyze/tests/test_fuse_matcher.py b/metainfer/tasks/sglang_trace_analyze/tests/test_fuse_matcher.py new file mode 100644 index 00000000..47fcd8f1 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/tests/test_fuse_matcher.py @@ -0,0 +1,37 @@ +"""Fuse matcher tests.""" + +from ..orchestrator.fuse_matcher import _match_consecutive, match_fuse_patterns + + +def test_match_consecutive_found(): + names = ["abc", "rms_norm", "triton_gemm", "add"] + pattern = ["rms_norm", "gemm"] + result = _match_consecutive(names, pattern) + assert result == ["rms_norm", "triton_gemm"] + + +def test_match_consecutive_not_found(): + names = ["abc", "rms_norm", "add"] + pattern = ["rms_norm", "gemm"] + result = _match_consecutive(names, pattern) + assert result == [] + + +def test_match_consecutive_short_list(): + names = ["abc"] + pattern = ["a", "b"] + result = _match_consecutive(names, pattern) + assert result == [] + + +def test_match_fuse_patterns_with_known_kernels(): + kernels = [ + {"kernel_name": "abc"}, + {"kernel_name": "triton_gemm"}, + {"kernel_name": "ncclAllReduce"}, + {"kernel_name": "triton_gemm"}, + ] + matches = match_fuse_patterns(kernels) + # The "nccl_allreduce + gemm (no overlap)" pattern should fire + pattern_names = [m["pattern"] for m in matches] + assert "nccl_allreduce + gemm (no overlap)" in pattern_names diff --git a/metainfer/tasks/sglang_trace_analyze/tests/test_gpu_specs.py b/metainfer/tasks/sglang_trace_analyze/tests/test_gpu_specs.py new file mode 100644 index 00000000..0c8cb5b0 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/tests/test_gpu_specs.py @@ -0,0 +1,17 @@ +"""GPU specs lookup tests.""" + +from ..orchestrator.gpu_specs import GPU_SPECS, GpuSpec + + +def test_gpu_specs_known(): + for label in ("K100", "A100_80G", "H100", "B200"): + spec = GPU_SPECS.get(label) + assert spec is not None, f"missing spec for {label}" + assert spec.bf16_tflops > 0 + assert spec.bandwidth_gb_s > 0 + + +def test_gpu_specs_values_reasonable(): + k100 = GPU_SPECS["K100"] + assert k100.bf16_tflops == 192 + assert k100.bandwidth_gb_s == 700 diff --git a/metainfer/tasks/sglang_trace_analyze/tests/test_plugin.py b/metainfer/tasks/sglang_trace_analyze/tests/test_plugin.py new file mode 100644 index 00000000..357b59d8 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/tests/test_plugin.py @@ -0,0 +1,18 @@ +"""Validate plugin registration and import sanity.""" + +from metainfer.server.registry import all_plugins +from metainfer.orchestrator.tasks import all_tasks + + +def test_all_plugins_includes_sglang_trace_analyze(): + types = [p.type for p in all_plugins()] + assert "sglang_trace_analyze" in types, ( + f"sglang_trace_analyze not found in registered plugins: {types}" + ) + + +def test_all_tasks_includes_sglang_trace_analyze(): + task_types = [p.task_type for p in all_tasks()] + assert "sglang_trace_analyze" in task_types, ( + f"sglang_trace_analyze not found in registered tasks: {task_types}" + ) diff --git a/metainfer/tasks/sglang_trace_analyze/tests/test_server_readers.py b/metainfer/tasks/sglang_trace_analyze/tests/test_server_readers.py new file mode 100644 index 00000000..385fe8f7 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/tests/test_server_readers.py @@ -0,0 +1,75 @@ +"""Server state reader tests.""" + +import json +import tempfile +from pathlib import Path + +from ..server._state_readers import ( + read_batch_detail, + read_hints, + read_mapping, + read_summary, +) + + +def test_read_summary(): + with tempfile.TemporaryDirectory() as td: + state_dir = Path(td) + analysis = state_dir / "analysis" + analysis.mkdir() + (analysis / "summary.json").write_text( + json.dumps({"model": "test", "batches": []}) + ) + result = read_summary(state_dir) + assert result is not None + assert result["model"] == "test" + + +def test_read_summary_missing(): + with tempfile.TemporaryDirectory() as td: + assert read_summary(Path(td)) is None + + +def test_read_mapping(): + with tempfile.TemporaryDirectory() as td: + state_dir = Path(td) + analysis = state_dir / "analysis" + analysis.mkdir() + (analysis / "mapping.json").write_text( + json.dumps({"entries": [{"kernel_name": "test"}]}) + ) + result = read_mapping(state_dir) + assert len(result["entries"]) == 1 + + +def test_read_hints(): + with tempfile.TemporaryDirectory() as td: + state_dir = Path(td) + analysis = state_dir / "analysis" + analysis.mkdir() + (analysis / "hints.json").write_text( + json.dumps({"bottleneck": {"kernel_or_pattern": "triton_gemm"}}) + ) + result = read_hints(state_dir) + assert result["bottleneck"]["kernel_or_pattern"] == "triton_gemm" + + +def test_read_batch_detail(): + with tempfile.TemporaryDirectory() as td: + state_dir = Path(td) + batch_dir = state_dir / "analysis" / "batches" / "bs_8" / "decode" + batch_dir.mkdir(parents=True) + (batch_dir / "kernel_table.json").write_text(json.dumps({"kernels": []})) + (batch_dir / "overlap.json").write_text(json.dumps({"gaps": []})) + (batch_dir / "fuse.json").write_text(json.dumps({"matches": []})) + + result = read_batch_detail(state_dir, 8, "decode") + assert result is not None + assert result["kernel_table"]["kernels"] == [] + assert result["overlap"]["gaps"] == [] + + +def test_read_batch_detail_missing(): + with tempfile.TemporaryDirectory() as td: + result = read_batch_detail(Path(td), 8, "decode") + assert result is None diff --git a/metainfer/tasks/sglang_trace_analyze/tests/test_structure_mapper.py b/metainfer/tasks/sglang_trace_analyze/tests/test_structure_mapper.py new file mode 100644 index 00000000..42ec91c5 --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/tests/test_structure_mapper.py @@ -0,0 +1,63 @@ +"""Structure mapper tests.""" + +from ..orchestrator.structure_mapper import ( + _infer_layer, + _infer_op_type, + build_mapping, +) + + +def test_infer_layer_from_call_stack(): + stack = " File \"sglang/srt/layers/attn/triton_ops.py\", line 45\n File \"model.py\"" + layer = _infer_layer(stack, "triton_attn_kernel", {}) + assert "attn" in layer.lower() if layer else True # matched sglang path + + +def test_infer_layer_model_layers_pattern(): + stack = "model.layers.5.self_attn.qkv_proj" + layer = _infer_layer(stack, "triton_gemm", {}) + assert layer == "layer_5" + + +def test_infer_op_type_attention(): + assert _infer_op_type("flash_attn_fwd", "") == "Attention" + assert _infer_op_type("flash_fwd_splitkv_mla", "") == "Attention" + + +def test_infer_op_type_gemm(): + assert _infer_op_type("triton_gemm_kernel", "") == "GEMM" + assert _infer_op_type("w8a8_bf16_matmul", "") == "GEMM" + + +def test_infer_op_type_moe(): + assert _infer_op_type("fused_moe_kernel", "") == "MoE" + + +def test_infer_op_type_norm(): + assert _infer_op_type("rms_norm_kernel", "") == "Norm" + + +def test_infer_op_type_nccl(): + assert _infer_op_type("ncclAllReduce", "") == "NCCL" + + +def test_build_mapping_empty(): + entries = build_mapping([], {"num_hidden_layers": 32}) + assert entries == [] + + +def test_build_mapping_with_call_stack(): + kernels = [ + { + "kernel_name": "triton_gemm", + "total_dur_us": 1000, + "count": 10, + "call_stack": "model.layers.3.self_attn.q_proj", + } + ] + entries = build_mapping(kernels, {"num_hidden_layers": 32}) + assert len(entries) == 1 + assert entries[0]["kernel_name"] == "triton_gemm" + assert entries[0]["model_layer"] == "layer_3" + assert entries[0]["op_type"] == "GEMM" + assert entries[0]["confidence"] == "high" diff --git a/metainfer/tasks/sglang_trace_analyze/tests/test_trace_parser.py b/metainfer/tasks/sglang_trace_analyze/tests/test_trace_parser.py new file mode 100644 index 00000000..c02b1afc --- /dev/null +++ b/metainfer/tasks/sglang_trace_analyze/tests/test_trace_parser.py @@ -0,0 +1,66 @@ +"""Trace parser tests with synthetic trace fixtures.""" + +import json +from ..orchestrator.trace_parser import aggregate_kernels + + +def _synthetic_trace(kernels): + """Build a minimal Chrome trace JSON.""" + events = [] + for name, dur, extra in kernels: + evt = {"cat": "kernel", "name": name, "ph": "X", "dur": dur, "ts": 0} + if extra: + evt.setdefault("args", {}).update(extra) + events.append(evt) + return {"traceEvents": events} + + +def test_aggregate_empty_trace(): + result = aggregate_kernels(_synthetic_trace([])) + assert result == [] + + +def test_aggregate_single_kernel(): + trace = _synthetic_trace([("triton_gemm", 1000, {})]) + result = aggregate_kernels(trace) + assert len(result) == 1 + assert result[0]["kernel_name"] == "triton_gemm" + assert result[0]["total_dur_us"] == 1000 + assert result[0]["count"] == 1 + + +def test_aggregate_multiple_same_kernel(): + trace = _synthetic_trace([ + ("triton_gemm", 500, {}), + ("triton_gemm", 700, {}), + ("flash_attn", 300, {}), + ]) + result = aggregate_kernels(trace) + assert len(result) == 2 + # triton_gemm aggregates: 500 + 700 = 1200 + assert result[0]["kernel_name"] == "triton_gemm" + assert result[0]["total_dur_us"] == 1200 + assert result[0]["count"] == 2 + # flash_attn is second + assert result[1]["kernel_name"] == "flash_attn" + assert result[1]["total_dur_us"] == 300 + + +def test_aggregate_ignores_non_kernel(): + trace = _synthetic_trace([ + ("triton_gemm", 500, {}), + ("cpu_op", 200, {}), # different cat + ]) + # Make the second event non-kernel + trace["traceEvents"][1]["cat"] = "cpu_op" + result = aggregate_kernels(trace) + assert len(result) == 1 + assert result[0]["kernel_name"] == "triton_gemm" + + +def test_aggregate_includes_call_stack(): + trace = _synthetic_trace([ + ("triton_gemm", 500, {"call stack": "model.layers.5.self_attn"}), + ]) + result = aggregate_kernels(trace, include_call_stack=True) + assert result[0]["call_stack"] == "model.layers.5.self_attn"