Fix forge bench-mode helper imports from workspace root

src/perf_helper_materialization.py

@@ -44,7 +44,7 @@
 )
 MARK_STARTS = (MARK_START, OLD_MARK_START, LEGACY_MARK_START)
 MARK_END = "# <<< AKA-GENERATED <<<"
-FUNC_ANCHOR = "def _measure_cuda_event_fallback("
+VLLM_BLOCK_ANCHOR = "import sys as _aka_sys"
 VLLM_HELPER_SYMBOLS = ("_measure_cuda_event_fallback", "_benchmark_cuda_graph_or_events")
 
 ROCMBENCH_HELPER_STUB = '''"""Generated at workspace setup from src/tools/perf/performance_utils_pytest.py.
@@ -129,7 +129,7 @@ def canonical_aka_helper(root: Path = ROOT) -> str:
 
 def canonical_vllm_block(root: Path = ROOT) -> str:
     text = (root / "src/tools/perf/vllm_cuda_graph_block.py").read_text()
-    return text[text.index(FUNC_ANCHOR) :]
+    return text[text.index(VLLM_BLOCK_ANCHOR) :]
 
 
 def replace_marked_region(current: str, block: str) -> str | None:

src/tools/perf/vllm_cuda_graph_block.py

@@ -4,6 +4,17 @@
 in the sibling materialized ``_aka_benchmark.py`` module.
 """
 
+import sys as _aka_sys
+from pathlib import Path as _AkaPath
+
+
+# Task-provided forge drivers are copied to the workspace root and load this
+# runner by file path.  Python does not add a file-loaded module's directory to
+# sys.path, so make the sibling materialized helper discoverable explicitly.
+_AKA_HELPER_DIR = str(_AkaPath(__file__).resolve().parent)
+if _AKA_HELPER_DIR not in _aka_sys.path:
+    _aka_sys.path.insert(0, _AKA_HELPER_DIR)
+
 
 def _measure_cuda_event_fallback(fn, repetition, prepare_fn=None):
     try:

tests/test_perf_helper_materialization.py

@@ -57,6 +57,45 @@ def test_materializes_vllm_adapter_and_sibling_helper(tmp_path):
     assert (scripts / AKA_HELPER_FILE_NAME).read_text() == canonical_aka_helper(ROOT)
 
 
+def test_file_loaded_vllm_runner_finds_sibling_helper_from_workspace_root(tmp_path):
+    scripts = tmp_path / "scripts"
+    scripts.mkdir()
+    runner = scripts / "task_runner.py"
+    runner.write_text(f"{MARK_START}\n{VLLM_HELPER_STUB_BLOCK}{MARK_END}\n")
+    (tmp_path / "config.yaml").write_text(
+        "performance_command:\n  - python3 scripts/task_runner.py performance\n"
+    )
+    materialize_perf_helpers_in_workspace(tmp_path)
+
+    # Match forge's invocation shape: forge_driver.py runs at the workspace
+    # root and imports scripts/task_runner.py with spec_from_file_location().
+    probe = """
+import importlib.util
+from pathlib import Path
+
+runner = Path("scripts/task_runner.py").resolve()
+scripts = str(runner.parent)
+assert scripts not in __import__("sys").path
+spec = importlib.util.spec_from_file_location("_forge_task_runner", runner)
+module = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(module)
+helper = importlib.util.find_spec("_aka_benchmark")
+assert helper is not None
+print(Path(helper.origin).resolve())
+"""
+    result = subprocess.run(
+        [sys.executable, "-S", "-c", probe],
+        cwd=tmp_path,
+        capture_output=True,
+        text=True,
+        check=False,
+        env={"PYTHONPATH": ""},
+    )
+
+    assert result.returncode == 0, result.stderr or result.stdout
+    assert Path(result.stdout.strip()) == (scripts / AKA_HELPER_FILE_NAME).resolve()
+
+
 def test_materializes_helper_beside_eval_tools_entrypoint(tmp_path):
     eval_tools = tmp_path / "eval_tools"
     eval_tools.mkdir()