feat(bypass): kernel source resolution + shape recovery for graph/Triton kernels

src/hyperloom/agents/kernel/tests/test_bypass_classify.py

@@ -0,0 +1,114 @@
+###############################################################################
+# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc.
+# SPDX-License-Identifier: MIT
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Unit tests for :mod:`_bypass_classify.classify_kernel` taxonomy coverage."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import pytest
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "tools"))
+
+from _bypass_classify import classify_kernel  # noqa: E402
+
+
+# (device_kernel_name, expected_category) covering real sglang+aiter kernels.
+_CATEGORY_CASES = [
+    # Communication / collective.
+    ("cross_device_reduce_2stage", "Communication"),
+    ("sglang::outplace_all_reduce", "Communication"),
+    ("all_reduce", "Communication"),
+    ("allgather_vec", "Communication"),
+    ("reg_all_gather", "Communication"),
+    ("_ZN5aiter26cross_device_reduce_2stageIDF16bLi8ELb0EEEvPNS_8RankDataE", "Communication"),
+    # Attention.
+    ("_fwd_grouped_kernel_stage1", "SDPA"),
+    ("_fwd_kernel_stage2", "SDPA"),
+    ("_decode_grouped_att", "SDPA"),
+    ("mla_decode", "SDPA"),
+    ("paged_attention", "SDPA"),
+    ("_score_kernel", "SDPA"),
+    ("_combine_kernel", "SDPA"),
+    ("kda_packed_decode", "SDPA"),
+    # MoE.
+    ("fused_moe", "MoE"),
+    ("moe_gemm1", "MoE"),
+    ("moe_gemm2", "MoE"),
+    ("ck_moe", "MoE"),
+    ("fmoe", "MoE"),
+    ("moe_sorting", "MoE"),
+    ("grouped_topk", "MoE"),
+    ("moe_align", "MoE"),
+    ("topk_softmax", "MoE"),
+    ("biased_grouped_topk", "MoE"),
+    # GEMM.
+    ("gemm_a16w16", "GEMM"),
+    ("gemm_a8w8", "GEMM"),
+    ("hgemm_bf16", "GEMM"),
+    ("flatmm", "GEMM"),
+    ("Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT160x64x128", "GEMM"),
+    ("opus_gemm", "GEMM"),
+    # Normalization.
+    ("add_rmsnorm", "Normalization"),
+    ("rmsnorm", "Normalization"),
+    ("layernorm", "Normalization"),
+    ("rms_norm", "Normalization"),
+    ("add_rmsnorm_quant_kernel", "Normalization"),
+    # Elementwise.
+    ("silu_and_mul", "Elementwise"),
+    ("situ_and_mul", "Elementwise"),
+    ("act_and_mul", "Elementwise"),
+    ("elementwise_kernel", "Elementwise"),
+    ("index_elementwise", "Elementwise"),
+    ("copy_", "Elementwise"),
+    ("_to_copy", "Elementwise"),
+    ("void (anonymous namespace)::situ_and_mul_kernel", "Elementwise"),
+    # Quantization.
+    ("quant", "Quantization"),
+    ("dynamic_per_token", "Quantization"),
+    ("per_tensor_quant", "Quantization"),
+    ("static_quant_fp8", "Quantization"),
+]
+
+
+@pytest.mark.parametrize("name,expected", _CATEGORY_CASES)
+def test_category_mapping(name: str, expected: str) -> None:
+    assert classify_kernel(name).category == expected, name
+
+
+def test_communication_is_not_reusable() -> None:
+    for name in [
+        "cross_device_reduce_2stage",
+        "sglang::outplace_all_reduce",
+        "all_reduce",
+        "allgather_vec",
+        "reg_all_gather",
+        "_ZN5aiter26cross_device_reduce_2stageIDF16bLi8ELb0EEEvPNS_8RankDataE",
+    ]:
+        kc = classify_kernel(name)
+        assert kc.category == "Communication", name
+        assert kc.reusable is False, name
+
+
+def test_mangled_allreduce_via_op_name_fallback() -> None:
+    # op_name attribution present: still Communication and not reusable.
+    kc = classify_kernel("_ZN5aiter26unknown_mangledEv", op_name="sglang::outplace_all_reduce")
+    assert kc.category == "Communication"
+    assert kc.reusable is False
+
+
+def test_moe_wins_over_gemm_substring() -> None:
+    assert classify_kernel("moe_gemm1").category == "MoE"
+
+
+def test_kda_packed_decode_is_reusable_sdpa() -> None:
+    kc = classify_kernel("void (anonymous namespace)::kda_packed_decode_kernel<8, false>(x)")
+    assert kc.category == "SDPA"
+    assert kc.reusable is True

src/hyperloom/agents/kernel/tests/test_bypass_reader.py

@@ -78,6 +78,36 @@ def test_correlation_attribution(tmp_path):
     assert kernels["paged_attention_v1"]["op_name"] == ""
 
 
+_FALLBACK_TRACE_EVENTS = [
+    # Capture-time launch of add_rmsnorm: cpu_op carries Input Dims (shape source).
+    {"cat": "cpu_op", "name": "aiter::add_rmsnorm", "args": {"External id": 100, "Input Dims": [[17, 7168]], "Input type": ["c10::BFloat16"]}},
+    {"cat": "cuda_runtime", "name": "hipLaunchKernel", "args": {"correlation": 5, "External id": 100}},
+    {"cat": "kernel", "ph": "X", "name": "add_rmsnorm_kernel", "ts": 1000, "dur": 100, "args": {"correlation": 5, "grid": [17, 1, 1], "block": [256, 1, 1]}},
+    # Graph-replay of the SAME kernel: no cpu_op link, grid-less Dispatch Task.
+    {"cat": "cuda_runtime", "name": "hipGraphLaunch", "args": {"correlation": 6}},
+    {"cat": "kernel", "ph": "X", "name": "add_rmsnorm_kernel", "ts": 2000, "dur": 100, "args": {"correlation": 6}},
+    # Pure-Triton kernel: no cpu_op, but launch carries grid/block geometry.
+    {"cat": "cuda_runtime", "name": "hipModuleLaunchKernel", "args": {"correlation": 7}},
+    {"cat": "kernel", "ph": "X", "name": "_score_kernel", "ts": 3000, "dur": 100, "args": {"correlation": 7, "grid": [17, 2, 1], "block": [512, 1, 1]}},
+]
+
+
+def test_launch_geom_and_backfill_threaded_to_rows(tmp_path):
+    payload = json.dumps({"traceEvents": _FALLBACK_TRACE_EVENTS}).encode("utf-8")
+    tf = tmp_path / "fb.trace.json"
+    tf.write_bytes(payload)
+    out = reader.analyze_trace(tf, top_k=0)
+    rows = {k["name"]: k for k in out["kernels"]}
+
+    # Same-name kernel resolved a shape at capture time -> backfill index seeded.
+    assert rows["add_rmsnorm_kernel"]["backfill_shapes"] == [[17, 7168]]
+    # Pure-Triton kernel: no cpu_op shape, but launch geometry is retained.
+    score = rows["_score_kernel"]
+    assert score["op_shapes"] == []
+    assert score["launch_grid"] == [17, 2, 1]
+    assert score["launch_block"] == [512, 1, 1]
+
+
 def test_timeline_union_math(tmp_path):
     tf = _write_trace(tmp_path / "t.trace.json")
     tl = reader.analyze_trace(tf, top_k=0)["timeline"]

src/hyperloom/agents/kernel/tests/test_bypass_report.py

@@ -61,6 +61,92 @@ def test_build_candidates_routing():
     assert any(c["name"] == "aten::mm" for c in cands["skipped_kernels"])
 
 
+def _one(cand_kernel):
+    """Build candidates for a single kernel row and return its candidate dict."""
+    cands = report.build_candidates(_analyze([cand_kernel]), framework="sglang", target_platform="MI300X")
+    return cands["hot_kernels"][0]
+
+
+def test_shape_provenance_torch_trace_wins():
+    # A kernel with its own cpu_op Input Dims resolves to torch_trace even when
+    # backfill/launch-grid data is also present (waterfall priority).
+    c = _one(
+        {
+            "name": "_score_kernel",
+            "op_name": "aten::score",
+            "gpu_time_us": 100.0,
+            "count": 1,
+            "op_shapes": [[17, 7168]],
+            "op_dtypes": ["c10::BFloat16"],
+            "backfill_shapes": [[99, 99]],
+            "launch_grid": [17, 2, 1],
+            "launch_block": [512, 1, 1],
+        }
+    )
+    assert c["shape_provenance"] == "torch_trace"
+    assert c["shapes"] and "17,7168" in c["shapes"][0]["shape"]
+
+
+def test_shape_provenance_capture_backfill():
+    # No own cpu_op shape, but the same-name kernel resolved a shape at capture
+    # time: inherit it, tagged capture_backfill.
+    c = _one(
+        {
+            "name": "aiter::add_rmsnorm_quant_kernel",
+            "op_name": "",
+            "gpu_time_us": 100.0,
+            "count": 1,
+            "op_shapes": [],
+            "backfill_shapes": [[17, 7168]],
+            "backfill_dtypes": ["c10::BFloat16"],
+            "launch_grid": [17, 1, 1],
+        }
+    )
+    assert c["shape_provenance"] == "capture_backfill"
+    assert c["shapes"] and "17,7168" in c["shapes"][0]["shape"]
+
+
+def test_shape_provenance_launch_grid():
+    # No cpu_op shape and no backfill: fall to launch grid/block geometry.
+    c = _one(
+        {
+            "name": "_combine_kernel",
+            "op_name": "",
+            "gpu_time_us": 100.0,
+            "count": 1,
+            "launch_grid": [17, 7, 1],
+            "launch_block": [256, 1, 1],
+        }
+    )
+    assert c["shape_provenance"] == "launch_grid"
+    s = c["shapes"][0]["shape"]
+    assert "grid=(17,7,1)" in s and "block=(256,1,1)" in s
+    # Geometry must NOT feed the analytical roofline.
+    assert c["roofline_source"] == report._RL_PLACEHOLDER
+
+
+def test_shape_provenance_tile_name():
+    # No shape and no launch geometry: parse the BLOCK_SIZE_* tile from the name.
+    c = _one(
+        {
+            "name": "_gemm_a16_w16_kernel_BLOCK_SIZE_M_32_BLOCK_SIZE_N_32_BLOCK_SIZE_K_256_x",
+            "op_name": "",
+            "gpu_time_us": 100.0,
+            "count": 1,
+        }
+    )
+    assert c["shape_provenance"] == "tile_name"
+    s = c["shapes"][0]["shape"]
+    assert "M32" in s and "N32" in s and "K256" in s
+
+
+def test_shape_provenance_unresolved_when_no_signal():
+    # Nothing to go on: stays unresolved (unchanged behaviour).
+    c = _one({"name": "mystery_kernel", "op_name": "", "gpu_time_us": 100.0, "count": 1})
+    assert c["shape_provenance"] == "unresolved"
+    assert c["shapes"] == []
+
+
 def test_build_workload_roofline_totals_covers_all_kernels():
     # 20 kernels > any top-k cap: the workload totals must sum every device
     # kernel, not just the top-k candidate list.

src/hyperloom/agents/kernel/tests/test_bypass_source_resolver.py

@@ -14,9 +14,12 @@
 
 from __future__ import annotations
 
+import shutil
 import sys
 from pathlib import Path
 
+import pytest
+
 sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "tools"))
 
 import _bypass_source_resolver as resolver  # noqa: E402
@@ -145,3 +148,53 @@ def test_editable_trace_source_rejects_generated_and_empty():
 def test_missing_json_yields_unresolved(monkeypatch):
     _patch_mapping(monkeypatch, {})
     assert resolver.resolve_source("anything", framework="vllm") == ("", "unresolved")
+
+
+def test_demangle_itanium_nested():
+    n = "_ZN5aiter26cross_device_reduce_2stageIDF16bLi8ELb0EEEvPNS_8RankDataE"
+    assert resolver._demangle_kernel_name(n) == "cross_device_reduce_2stage"
+
+
+def test_demangle_plain_triton_name_passthrough():
+    assert resolver._demangle_kernel_name("_fwd_grouped_kernel_stage1") == "_fwd_grouped_kernel_stage1"
+
+
+def test_demangle_anonymous_namespace_template():
+    n = "void (anonymous namespace)::kda_packed_decode_kernel<8, false>(x)"
+    assert resolver._demangle_kernel_name(n) == "kda_packed_decode_kernel"
+
+
+@pytest.fixture
+def repo_dir():
+    """A repo-like dir avoiding /tmp and the 'test' skip marker in its path."""
+    base = Path(__file__).resolve().parents[2] / "_bypass_repo_scan_fixture" / "src"
+    base.mkdir(parents=True, exist_ok=True)
+    yield base
+    shutil.rmtree(base.parent, ignore_errors=True)
+
+
+def test_resolve_by_kernel_name_triton_and_native(monkeypatch, repo_dir):
+    py = repo_dir / "fused.py"
+    py.write_text("@triton.jit\ndef foo(x):\n    return x\n", encoding="utf-8")
+    cu = repo_dir / "kern.cu"
+    cu.write_text("__global__ void bar(float* p) {}\n", encoding="utf-8")
+    monkeypatch.setattr(
+        resolver,
+        "_build_repo_kernel_index",
+        lambda: {"foo": str(py), "bar": str(cu), "gen": "/tmp/torchinductor_x/gen.py"},
+    )
+    assert resolver.resolve_by_kernel_name("foo") == (str(py), "repo_scan")
+    assert resolver.resolve_by_kernel_name("bar") == (str(cu), "repo_scan")
+    assert resolver.resolve_by_kernel_name("gen") == ("", "unresolved")
+    assert resolver.resolve_by_kernel_name("missing") == ("", "unresolved")
+
+
+def test_build_repo_kernel_index_scans_roots(monkeypatch, repo_dir):
+    (repo_dir / "a.py").write_text("@triton.jit\ndef tri_k(x):\n    pass\n", encoding="utf-8")
+    (repo_dir / "b.cu").write_text("__global__ void nat_k(int* p) {}\n", encoding="utf-8")
+    monkeypatch.setattr(resolver, "_repo_scan_roots", lambda: (str(repo_dir),))
+    resolver._build_repo_kernel_index.cache_clear()
+    index = resolver._build_repo_kernel_index()
+    resolver._build_repo_kernel_index.cache_clear()
+    assert index["tri_k"] == str(repo_dir / "a.py")
+    assert index["nat_k"] == str(repo_dir / "b.cu")

src/hyperloom/agents/kernel/tests/test_bypass_trace_analysis.py

@@ -791,3 +791,70 @@ def test_discover_capture_shards_dedups_tp_ranks(tmp_path):
     labels = sorted(lbl for _p, lbl, _m in shards)
     assert labels == ["bs_16", "bs_64"]  # bs_16 deduped 2 ranks -> 1
     assert len(shards) == 2
+
+
+# ── CUDA/HIP graph-mode under-recording ──────────────────────────────────────
+
+
+def _graph_under_recorded_events():
+    """Synthetic graph-mode trace: 4 graph launches but only 1 replay's kernels
+    recorded, spanning a large wall clock (busy fraction << 0.5)."""
+    events = [{"cat": "cpu_op", "name": "aten::mm", "args": {"External id": 200}}]
+    # Four graph-launch runtime events (no External id, only correlation).
+    for i, corr in enumerate((5, 6, 7, 8)):
+        events.append(
+            {"cat": "cuda_runtime", "name": "hipGraphLaunch", "args": {"correlation": corr, "External id": None}}
+        )
+    # A normally-launched kernel (attributable) and graph-internal kernels for a
+    # single recorded replay (correlation 5). Total wall span ~10s, busy ~200us.
+    events += [
+        {"cat": "cuda_runtime", "name": "hipLaunchKernel", "args": {"correlation": 99, "External id": 200}},
+        {"cat": "kernel", "ph": "X", "name": "Cijk_Alik_Bljk_HHS", "ts": 1000, "dur": 100, "args": {"correlation": 99}},
+        {"cat": "kernel", "ph": "X", "name": "graph_fused_attn", "ts": 1100, "dur": 100, "args": {"correlation": 5}},
+        {"cat": "kernel", "ph": "X", "name": "graph_fused_mlp", "ts": 10_000_000, "dur": 100, "args": {"correlation": 5}},
+    ]
+    return events
+
+
+def test_reader_marks_graph_attributed_not_unlinked(tmp_path):
+    # Graph-launch correlations classify replayed kernels as graph-attributed,
+    # never as (unlinked); flags graph_mode / graph_under_recorded.
+    trace = tmp_path / "g.trace.json"
+    trace.write_bytes(json.dumps({"traceEvents": _graph_under_recorded_events()}).encode("utf-8"))
+    out = bta._reader.analyze_trace(str(trace), top_k=0)
+    attr = out["attribution"]
+    cov = out["graph_coverage"]
+    assert attr["graph_mode"] is True
+    assert attr["graph_launch_count"] == 4
+    assert attr["graph_attributed_kernels"] == 2
+    assert attr["unlinked_kernels"] == 0
+    assert cov["graph_under_recorded"] is True
+    assert cov["busy_fraction"] < 0.5
+    # graph-internal kernels are not surfaced as a (unlinked) op bucket
+    assert "(unlinked)" not in {o["name"] for o in out["ops"]}
+
+
+def test_reader_no_graph_mode_when_no_graph_launches(tmp_path):
+    trace = tmp_path / "ng.trace.json"
+    trace.write_bytes(json.dumps({"traceEvents": _TRACE_EVENTS}).encode("utf-8"))
+    out = bta._reader.analyze_trace(str(trace), top_k=0)
+    assert out["attribution"]["graph_mode"] is False
+    assert out["graph_coverage"]["graph_under_recorded"] is False
+
+
+def test_graph_under_recorded_skips_idle_gate_keeps_candidates(tmp_path, capsys, monkeypatch):
+    # Under-recorded graph trace: idle% is ~99% but the idle gate must NOT clear
+    # candidates; instead a bypass_graph_under_recorded warning is surfaced.
+    monkeypatch.delenv("HYPERLOOM_BYPASS_STEADY_STATE", raising=False)
+    monkeypatch.delenv("HYPERLOOM_TRACELENS_IDLE_PCT_THRESHOLD", raising=False)
+    trace = tmp_path / "g.trace.json"
+    trace.write_bytes(json.dumps({"traceEvents": _graph_under_recorded_events()}).encode("utf-8"))
+    rc, result, _ = _run(_base_argv(tmp_path, str(trace)), capsys)
+    assert rc == 0
+    assert result["timeline"]["idle_pct"] > 80.0
+    assert result["graph_coverage"]["graph_under_recorded"] is True
+    codes = {w["code"] for w in result["trace_health_warnings"]}
+    assert "bypass_graph_under_recorded" in codes
+    assert "high_gpu_idle_pct" not in codes
+    # candidates are preserved (ranked by recorded-kernel GPU share)
+    assert result["hot_kernels"], "candidates must survive the idle gate under graph under-recording"