diff --git a/.github/workflows/radeon-forge.yml b/.github/workflows/radeon-forge.yml new file mode 100644 index 0000000000000..d2d02453ad2ab --- /dev/null +++ b/.github/workflows/radeon-forge.yml @@ -0,0 +1,39 @@ +name: Radeon Forge + +on: + push: + branches: [radeon-forge] + paths: + - "extra/radeon_forge/**" + - "test/test_radeon_forge*.py" + - ".github/workflows/radeon-forge.yml" + pull_request: + paths: + - "extra/radeon_forge/**" + - "test/test_radeon_forge*.py" + +permissions: + contents: read + +jobs: + forge: + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Compile Radeon Forge + run: python -m compileall -q extra/radeon_forge test/test_radeon_forge*.py + - name: Run isolated Forge tests + run: >- + python -m unittest -v + test.test_radeon_forge + test.test_radeon_forge_recipe + test.test_radeon_forge_runtime + test.test_radeon_forge_hooks + test.test_radeon_forge_autotune + test.test_radeon_forge_kv + test.test_radeon_forge_tool_stream + test.test_radeon_forge_openai diff --git a/extra/radeon_forge/KERNEL_CORPUS.md b/extra/radeon_forge/KERNEL_CORPUS.md new file mode 100644 index 0000000000000..1e44c187697a7 --- /dev/null +++ b/extra/radeon_forge/KERNEL_CORPUS.md @@ -0,0 +1,104 @@ +# Kernel and optimization corpus + +This file records the provenance and intended use of initial Radeon Forge implementation families. Inclusion does not imply that a kernel is correct or fast on the W7900. Every candidate must still pass target compilation, resource, numerical, benchmark, and held-out gates. + +## Native RDNA3 / gfx1100 seed + +### `extra/gemm/amd_asm_matmul.py` + +Status: first hardware family. + +Properties: + +- direct RDNA3 instruction construction through tinygrad's AMD DSL; +- explicitly targets gfx1100; +- 128x128 float32 GEMM tile; +- 128-thread workgroup; +- VOPD-aware accumulator and operand register placement; +- LDS staging and double-buffer-style global prefetch; +- reference check against tinygrad matmul already exists. + +Initial tuning/synthesis dimensions: + +- FMAC pair order and instruction scheduling; +- prefetch placement and distance; +- `s_clause` grouping; +- waitcnt placement; +- LDS swizzle and padding; +- occupancy limiter used by the harness; +- matrix-shape specialization. + +Unsafe dimensions are not exposed as blind scalar knobs. Structural variants must regenerate a complete program and pass the reference oracle. + +## Llama fused-kernel patterns + +Most current files under `extra/llama_kernels` compile with `HIPCCCompiler("gfx950", ...)`. They are therefore architectural patterns, not accepted gfx1100 candidates. + +### Fused RMSNorm / multiply / FP8 quantization + +Source: + +- `extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/` + +Useful ideas: + +- remove intermediate HBM materializations; +- one workgroup per row with grid-stride row processing; +- vectorized BF16 loads; +- fused normalization, weighting, quantization, saved backward state, and amax; +- compile-time workgroup and grid parameters. + +Candidate structural variants: + +- per-workgroup amax followed by a second reduction; +- scalar global atomic amax; +- workgroup/thread-count choices valid for hidden dimension; +- residual-add fusion; +- direct C/HIP versus UOp representation. + +Required before use on W7900: + +- successful gfx1100 compilation; +- confirmation that the relevant dtype/intrinsics are supported and performant; +- numerical validation against tinygrad; +- resource metadata and no-spill gate; +- end-to-end relevance to the selected inference model. + +### Fused cross entropy + +Provenance pattern: upstream tinygrad PR #16263 replaced handwritten C with a portable UOp version while preserving fusion and memory savings. The PR notes that the then-current BEAM-tuned UOp forward remained much slower than handwritten C, while end-to-end training differed far less. + +Use in Forge: + +- compare direct target-specific code and portable UOp families; +- optimize the end-to-end objective, not an isolated kernel ratio; +- retain the portable family as fallback. + +### Atomic amax + +Provenance pattern: upstream tinygrad PR #17063 replaces a two-stage amax reduction with atomics to remove repeated reduction-kernel launches. + +Use in Forge: + +- structural alternative, not assumed winner; +- compare under representative tensor size and contention; +- reject if numerical or end-to-end behavior regresses. + +## Compiler-resource oracle + +Provenance pattern: upstream tinygrad PR #3641 demonstrates extraction of COMGR executable metadata including VGPR/SGPR counts, LDS, scratch, occupancy, and spill counts. + +Forge uses these as hard or diagnostic signals: + +- spills may be forbidden by workload contract; +- excess VGPR/LDS use can reject a candidate before expensive trials; +- a kernel that loses occupancy is not automatically rejected, but the trade-off is measured and retained in evidence. + +## Corpus policy + +1. Record source and target architecture. +2. Distinguish a reusable idea from code verified on gfx1100. +3. Preserve rejected variants and reasons. +4. Never import a claimed benchmark without reproducing it on the assigned W7900. +5. Prefer end-to-end inference improvement over isolated-kernel speedup. +6. Keep tinygrad/UOp implementations as reference and fallback even when direct generated code wins. diff --git a/extra/radeon_forge/README.md b/extra/radeon_forge/README.md new file mode 100644 index 0000000000000..45cc01acebc01 --- /dev/null +++ b/extra/radeon_forge/README.md @@ -0,0 +1,160 @@ +# Radeon Forge + +Radeon Forge is an oracle-guided, fully local performance-engineering agent for private AI workloads on AMD Radeon / ROCm. + +The durable project is the workload intent, invariants, correctness oracle, benchmark protocol, hardware evidence, and experiment history. Generated kernel implementations are disposable candidates: Forge may regenerate or replace them, but it may not deploy one that fails a hard oracle. + +## Track 2 scenario + +A user asks Forge to optimize a locally deployed private agent under explicit constraints, for example: + +```text +Minimize P95 end-to-end task latency on a Radeon PRO W7900. +Do not use remote model APIs. Preserve task success within 1% of baseline, +require valid tool calls, forbid kernel spills, and stay below the VRAM limit. +``` + +Forge decomposes the request, invokes compiler/test/benchmark tools, remembers prior experiments, requests permission for consequential actions, and deploys or rolls back a measured configuration. + +## Optimization loop + +```text +contract + -> baseline + -> choose kernel/runtime family + -> compile for gfx1100 + -> reject invalid resource usage + -> correctness oracle + -> short hardware trials + -> successive halving + -> held-out validation + -> approval + -> deploy or revert +``` + +The agent may propose structural rewrites. The deterministic tuner evaluates parameters within each implementation family. Neither may override a failed correctness, quality, stability, privacy, or resource gate. + +## Metrics + +Primary submission metric: + +- P95 end-to-end latency on a frozen private-agent task suite. + +Supporting metrics: + +- time to first token; +- inter-token latency and decode throughput; +- kernel/subgraph median and P95 latency; +- task success and tool-call validity; +- numerical maximum absolute and relative error; +- VGPR, SGPR, LDS, scratch, spills, occupancy, and peak VRAM; +- crash rate and repeated-run stability; +- external network calls. + +Speed is optimized subject to correctness and quality constraints. A faster invalid candidate cannot win. + +## Tinygrad's role + +Tinygrad is used as: + +1. the trusted reference implementation and end-to-end workload; +2. the baseline and fallback runtime; +3. a corpus of AMD kernel implementations and optimization patterns; +4. the low-level AMD execution path for selected generated candidates. + +Radeon Forge is kept under `extra/radeon_forge` instead of being embedded deeply into tinygrad's compiler. This makes the oracle and tuning loop independently auditable and lets generated HIP/UOp/AMD-DSL implementations remain replaceable. + +## Current implementation + +Implemented: + +- workload, objective, correctness, and resource contracts; +- hard feasibility gates; +- append-only JSONL experiment ledger; +- successive-halving tuner; +- AMD metadata parsing for registers, LDS, scratch, occupancy, and spills; +- target-specific HIP compilation backend; +- executable JSON benchmark protocol; +- kernel-family search-space representation; +- native gfx1100 assembly-GEMM family with valid alternate FMAC schedules and occupancy limits; +- gfx1100 hardware workload adapter with a tinygrad numerical oracle; +- fused RMSNorm/multiply/FP8 pattern family, gated as unverified on gfx1100; +- deterministic local retrieval with file-and-line citations; +- loopback-only OpenAI-compatible planner client with no remote fallback; +- deny-by-default, scoped, expiring permission grants; +- multi-turn proposal, revision, approval, authorization, completion, and failure state machine; +- unit tests for metadata, hard gates, command protocol, permissions, local endpoint restrictions, retrieval, and multi-turn approval. + +Not yet claimed or implemented: + +- no W7900 performance result has been recorded; +- the gfx950 Llama kernels are pattern references until individually compiled and validated on gfx1100; +- no generated candidate is approved for deployment; +- no local planner model/runtime has yet been selected and benchmarked on the W7900; +- the interactive terminal/web surface is still pending; +- the frozen end-to-end private-agent task suite and held-out evaluation are still pending. + +## First gfx1100 hardware gate + +Review the candidate plan without running hardware: + +```bash +python3 -m extra.radeon_forge.cli --root . --family rdna3-asm-matmul --n 1024 --dry-run +``` + +After reviewing the printed plan, explicitly approve the local benchmark: + +```bash +python3 -m extra.radeon_forge.cli \ + --root . \ + --family rdna3-asm-matmul \ + --n 1024 \ + --budgets 3,10,30 \ + --approve-benchmark +``` + +This family validates the Forge search, permission, oracle, and evidence path. It is not the final Track 2 submission metric. + +## Benchmark executable protocol + +A workload harness receives: + +- `RADEON_FORGE_CANDIDATE_JSON` +- `RADEON_FORGE_BUDGET` + +Its final non-empty stdout line must be a JSON object such as: + +```json +{ + "compile_ok": true, + "stable": true, + "samples_us": [8.4, 8.2, 8.3], + "correctness": { + "passed": true, + "max_abs_error": 0.0002, + "max_rel_error": 0.0008, + "checked_values": 65536 + }, + "resources": { + "vgprs": 72, + "sgprs": 32, + "lds_bytes": 8192, + "scratch_bytes": 0, + "spilled_vgprs": 0, + "spilled_sgprs": 0 + }, + "evidence": { + "target": "gfx1100", + "reference": "tinygrad" + } +} +``` + +## Development validation + +```bash +python3 -m unittest test.test_radeon_forge -v +python3 -m compileall -q extra/radeon_forge +``` + +Hardware evidence must be generated on the assigned Radeon Cloud W7900 before any speed claim is made. diff --git a/extra/radeon_forge/__init__.py b/extra/radeon_forge/__init__.py new file mode 100644 index 0000000000000..198fa149c96ed --- /dev/null +++ b/extra/radeon_forge/__init__.py @@ -0,0 +1,7 @@ +"""Radeon Forge: oracle-guided synthesis and autotuning for AMD inference workloads.""" + +from .contracts import Candidate, Objective, TrialResult, WorkloadContract +from .ledger import ExperimentLedger +from .tuner import successive_halving + +__all__ = ["Candidate", "ExperimentLedger", "Objective", "TrialResult", "WorkloadContract", "successive_halving"] diff --git a/extra/radeon_forge/agent.py b/extra/radeon_forge/agent.py new file mode 100644 index 0000000000000..edcdb6a7d146d --- /dev/null +++ b/extra/radeon_forge/agent.py @@ -0,0 +1,139 @@ +from __future__ import annotations + +import secrets +from dataclasses import asdict, dataclass, replace +from enum import Enum +from typing import Any, Mapping, Protocol, Sequence + +from .contracts import WorkloadContract +from .knowledge import LocalKnowledgeBase, RetrievalHit +from .ledger import ExperimentLedger +from .permissions import Action, PermissionController, PermissionGrant +from .planner import PlannerDecision + + +class AgentState(str, Enum): + READY = "ready" + AWAITING_APPROVAL = "awaiting_approval" + APPROVED = "approved" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + + +class AgentStateError(RuntimeError): pass + + +class PlannerLike(Protocol): + def plan(self, request: str, contract: WorkloadContract, evidence: Sequence[RetrievalHit] = (), + memory: Sequence[Mapping[str, Any]] = ()) -> PlannerDecision: ... + + +@dataclass(frozen=True) +class Proposal: + proposal_id: str + user_request: str + decision: PlannerDecision + evidence_citations: tuple[str, ...] + contract: WorkloadContract + supersedes: str | None = None + + +@dataclass(frozen=True) +class ApprovalReceipt: + proposal_id: str + approved_actions: tuple[Action, ...] + user_reason: str + grant: PermissionGrant + + +class ForgeAgent: + """Stateful multi-turn coordinator for planning and permissioned execution. + + This object does not execute kernels by itself. It coordinates the local + planner, retrieval, memory, and user approval boundary. Deterministic tools + must call `authorize` immediately before each consequential action. + """ + + def __init__(self, contract: WorkloadContract, planner: PlannerLike, knowledge: LocalKnowledgeBase, + ledger: ExperimentLedger, permissions: PermissionController | None = None): + self.contract = contract + self.planner = planner + self.knowledge = knowledge + self.ledger = ledger + self.permissions = permissions or PermissionController() + self.state = AgentState.READY + self.current_proposal: Proposal | None = None + self.approval: ApprovalReceipt | None = None + self._conversation: list[dict[str, Any]] = [] + + def _memory(self, limit: int = 20) -> tuple[Mapping[str, Any], ...]: + records = tuple(self.ledger.records()) + return tuple(records[-limit:]) + + def propose(self, user_request: str, top_k: int = 5) -> Proposal: + if self.state is AgentState.RUNNING: raise AgentStateError("cannot create a new proposal while tools are running") + request = user_request.strip() + if not request: raise ValueError("user_request must not be empty") + hits = self.knowledge.search(request, top_k=top_k) + decision = self.planner.plan(request, self.contract, evidence=hits, memory=self._memory()) + previous = self.current_proposal.proposal_id if self.current_proposal is not None else None + proposal = Proposal(secrets.token_urlsafe(12), request, decision, tuple(hit.citation for hit in hits), self.contract, previous) + self.current_proposal = proposal + self.approval = None + self.state = AgentState.AWAITING_APPROVAL + self._conversation.append({"role": "user", "content": request}) + self._conversation.append({"role": "agent", "proposal_id": proposal.proposal_id, "decision": asdict(decision)}) + self.ledger.append("proposal_created", proposal) + return proposal + + def revise(self, user_request: str, contract_changes: Mapping[str, Any] | None = None) -> Proposal: + if contract_changes: + allowed = {field for field in self.contract.__dataclass_fields__ if field not in {"name", "target"}} + unknown = set(contract_changes) - allowed + if unknown: raise ValueError(f"unsupported contract fields: {sorted(unknown)}") + self.contract = replace(self.contract, **dict(contract_changes)) + self.ledger.append("contract_revised", {"changes": dict(contract_changes), "contract": self.contract}) + return self.propose(user_request) + + def approve(self, proposal_id: str, user_reason: str, max_tool_uses: int | None = None) -> ApprovalReceipt: + if self.state is not AgentState.AWAITING_APPROVAL or self.current_proposal is None: + raise AgentStateError("there is no proposal awaiting approval") + if proposal_id != self.current_proposal.proposal_id: raise AgentStateError("approval does not match the current proposal") + actions = self.current_proposal.decision.proposed_actions + if not actions: raise AgentStateError("proposal contains no executable actions") + reason = user_reason.strip() + if not reason: raise ValueError("approval requires a user reason") + uses = max_tool_uses if max_tool_uses is not None else max(1, self.current_proposal.decision.benchmark_budget + len(actions)) + grant = self.permissions.issue(actions, reason, max_uses=uses) + receipt = ApprovalReceipt(proposal_id, actions, reason, grant) + self.approval = receipt + self.state = AgentState.APPROVED + # Never persist the bearer token. The receipt fields excluding the grant are + # enough to audit what the user approved. + self.ledger.append("proposal_approved", {"proposal_id": proposal_id, "approved_actions": [action.value for action in actions], + "user_reason": reason, "max_tool_uses": uses}) + return receipt + + def authorize(self, action: Action) -> PermissionGrant: + if self.state not in {AgentState.APPROVED, AgentState.RUNNING} or self.approval is None: + raise AgentStateError(f"action {action.value} is not approved") + grant = self.permissions.authorize(self.approval.grant.token, action) + self.state = AgentState.RUNNING + self.ledger.append("tool_authorized", {"proposal_id": self.approval.proposal_id, "action": action.value, "reason": grant.reason}) + return grant + + def complete(self, result: Mapping[str, Any]) -> None: + if self.state not in {AgentState.APPROVED, AgentState.RUNNING}: raise AgentStateError("no approved execution is active") + self.state = AgentState.COMPLETED + self.ledger.append("execution_completed", {"proposal_id": self.approval.proposal_id if self.approval else None, "result": dict(result)}) + + def fail(self, reason: str, evidence: Mapping[str, Any] | None = None) -> None: + if self.state not in {AgentState.APPROVED, AgentState.RUNNING}: raise AgentStateError("no approved execution is active") + self.state = AgentState.FAILED + self.ledger.append("execution_failed", {"proposal_id": self.approval.proposal_id if self.approval else None, + "reason": reason, "evidence": dict(evidence or {})}) + + @property + def conversation(self) -> tuple[Mapping[str, Any], ...]: + return tuple(self._conversation) diff --git a/extra/radeon_forge/amd_metadata.py b/extra/radeon_forge/amd_metadata.py new file mode 100644 index 0000000000000..28999dfd90ad3 --- /dev/null +++ b/extra/radeon_forge/amd_metadata.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +import re +from collections.abc import Mapping +from typing import Any + +from .contracts import ResourceReport + + +_PATTERNS: dict[str, tuple[str, ...]] = { + "vgprs": (r"\bNumVgprs:\s*(\d+)", r"\.vgpr_count:\s*(\d+)", r"\bvgpr_count:\s*(\d+)"), + "sgprs": (r"\bNumSgprs:\s*(\d+)", r"\.sgpr_count:\s*(\d+)", r"\bsgpr_count:\s*(\d+)"), + "scratch_bytes": (r"\bScratchSize:\s*(\d+)", r"\.private_segment_fixed_size:\s*(\d+)", r"\bprivate_segment_fixed_size:\s*(\d+)"), + "lds_bytes": (r"\bLDSByteSize:\s*(\d+)", r"\.group_segment_fixed_size:\s*(\d+)", r"\bgroup_segment_fixed_size:\s*(\d+)"), + "occupancy": (r"\bOccupancy:\s*(\d+)",), + "spilled_vgprs": (r"\.vgpr_spill_count:\s*(\d+)", r"\bvgpr_spill_count:\s*(\d+)"), + "spilled_sgprs": (r"\.sgpr_spill_count:\s*(\d+)", r"\bsgpr_spill_count:\s*(\d+)"), +} + + +def _first_int(text: str, patterns: tuple[str, ...]) -> int | None: + for pattern in patterns: + if match := re.search(pattern, text, flags=re.IGNORECASE): return int(match.group(1)) + return None + + +def resource_report_from_text(text: str) -> ResourceReport: + values = {name: _first_int(text, patterns) for name, patterns in _PATTERNS.items()} + return ResourceReport( + vgprs=values["vgprs"], sgprs=values["sgprs"], lds_bytes=values["lds_bytes"], scratch_bytes=values["scratch_bytes"], + occupancy=values["occupancy"], spilled_vgprs=values["spilled_vgprs"] or 0, spilled_sgprs=values["spilled_sgprs"] or 0, + ) + + +def _as_int(value: Any) -> int | None: + if value is None: return None + try: return int(value) + except (TypeError, ValueError): return None + + +def resource_report_from_comgr(metadata: Mapping[str, Any], kernel_index: int = 0) -> ResourceReport: + """Parse COMGR executable metadata such as the structure exposed in tinygrad PR #3641.""" + kernels = metadata.get("amdhsa.kernels") + if not isinstance(kernels, list) or kernel_index >= len(kernels) or not isinstance(kernels[kernel_index], Mapping): + raise ValueError("COMGR metadata does not contain the requested amdhsa kernel") + kernel = kernels[kernel_index] + return ResourceReport( + vgprs=_as_int(kernel.get(".vgpr_count", kernel.get("vgpr_count"))), + sgprs=_as_int(kernel.get(".sgpr_count", kernel.get("sgpr_count"))), + lds_bytes=_as_int(kernel.get(".group_segment_fixed_size", kernel.get("group_segment_fixed_size"))), + scratch_bytes=_as_int(kernel.get(".private_segment_fixed_size", kernel.get("private_segment_fixed_size"))), + spilled_vgprs=_as_int(kernel.get(".vgpr_spill_count", kernel.get("vgpr_spill_count"))) or 0, + spilled_sgprs=_as_int(kernel.get(".sgpr_spill_count", kernel.get("sgpr_spill_count"))) or 0, + ) diff --git a/extra/radeon_forge/backends/cancellation.py b/extra/radeon_forge/backends/cancellation.py new file mode 100644 index 0000000000000..06686ffde3f29 --- /dev/null +++ b/extra/radeon_forge/backends/cancellation.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +import signal, threading + + +class SignalCancellation: + """Process-local cooperative cancellation set by SIGUSR1.""" + def __init__(self): + self._event = threading.Event() + self.installed = False + + def install(self) -> bool: + if not hasattr(signal, "SIGUSR1"): return False + signal.signal(signal.SIGUSR1, lambda _signum, _frame: self._event.set()) + self.installed = True + return True + + def reset(self) -> None: self._event.clear() + def request(self) -> None: self._event.set() + @property + def requested(self) -> bool: return self._event.is_set() diff --git a/extra/radeon_forge/backends/fake.py b/extra/radeon_forge/backends/fake.py new file mode 100644 index 0000000000000..d16196f38631f --- /dev/null +++ b/extra/radeon_forge/backends/fake.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +from typing import Iterator, Mapping, Any +from ..runtime.backend import BackendCapabilities, GenerationEvent, GenerationRequest + + +class ScriptedBackend: + """Deterministic backend for UI development and runtime tests without a GPU.""" + def __init__(self, responses: list[str] | None = None): self.responses = responses or ["Radeon Forge is ready."] + @property + def name(self) -> str: return "scripted-local" + @property + def capabilities(self) -> BackendCapabilities: return BackendCapabilities(kernel_metrics=True) + @property + def runtime_metadata(self) -> Mapping[str, Any]: + return {"runtime": "scripted", "architecture": "", "gpu": "", "model_family": "test"} + def stream(self, request: GenerationRequest) -> Iterator[GenerationEvent]: + text = self.responses.pop(0) if self.responses else "Done." + yield GenerationEvent("prefill", metrics={"wall_ms": 4.0, "prompt_tokens": 32, "prefix_reused_tokens": 24}) + parts = text.split(" ") + for index, token in enumerate(parts): + yield GenerationEvent("token", token + (" " if index < len(parts) - 1 else ""), + metrics={"index": index, "wall_ms": 2.0, "gpu_ms": 1.4, "kernel_count": 24}) + yield GenerationEvent("done", finish_reason="stop", metrics={"generated_tokens": len(parts)}) + def close(self) -> None: pass diff --git a/extra/radeon_forge/backends/kv_state.py b/extra/radeon_forge/backends/kv_state.py new file mode 100644 index 0000000000000..0b4a1a2adcc08 --- /dev/null +++ b/extra/radeon_forge/backends/kv_state.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Sequence + + +@dataclass +class KVReuseLedger: + """Tracks the only prefix that is safe to reuse from one resident KV cache. + + Generated output is not reusable until it has been fed back through the model + as a decode input. This distinction matters at length stops and tool + boundaries: emitted text and materialized KV state are not always identical. + """ + active_session: str | None = None + cached_tokens: list[int] = field(default_factory=list) + + def begin(self, session_id: str, prompt: Sequence[int]) -> int: + if self.active_session != session_id: + self.active_session = session_id + self.cached_tokens.clear() + common = 0 + for old, new in zip(self.cached_tokens, prompt): + if old != new: break + common += 1 + return common + + def commit_prefill(self, prompt_without_pending_decode_token: Sequence[int]) -> None: + self.cached_tokens = list(prompt_without_pending_decode_token) + + def commit_decode_input(self, token: int, expected_position: int) -> None: + if len(self.cached_tokens) != expected_position: + raise RuntimeError(f"KV ledger position mismatch: cached={len(self.cached_tokens)} expected={expected_position}") + self.cached_tokens.append(int(token)) + + def reset(self) -> None: + self.active_session = None + self.cached_tokens.clear() diff --git a/extra/radeon_forge/backends/plugins.py b/extra/radeon_forge/backends/plugins.py new file mode 100644 index 0000000000000..9bd9950cd0edf --- /dev/null +++ b/extra/radeon_forge/backends/plugins.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +import sys +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Mapping, Sequence + + +@dataclass(frozen=True) +class WorkerPlugin: + name: str + module: str + description: str + required: tuple[str, ...] + optional: tuple[str, ...] = () + defaults: Mapping[str, Any] = field(default_factory=dict) + + def command(self, values: Mapping[str, Any]) -> list[str]: + config={**dict(self.defaults),**{str(k):v for k,v in values.items() if v is not None}} + missing=[name for name in self.required if config.get(name) in (None,"")] + if missing: raise ValueError(f"worker plugin {self.name!r} is missing required fields: {missing}") + allowed=set(self.required)|set(self.optional) + unknown=sorted(set(config)-allowed) + if unknown: raise ValueError(f"worker plugin {self.name!r} received unsupported fields: {unknown}") + command=[sys.executable,"-m",self.module] + for name in self.required+self.optional: + if name not in config or config[name] is None: continue + value=config[name] + flag="--"+name.replace("_","-") + if isinstance(value,bool): + if value: command.append(flag) + elif isinstance(value,(str,int,float,Path)): command += [flag,str(value)] + elif isinstance(value,Sequence) and not isinstance(value,(str,bytes,bytearray)): + for item in value: command += [flag,str(item)] + else: raise TypeError(f"unsupported worker option {name}: {type(value).__name__}") + return command + + +class WorkerPluginRegistry: + def __init__(self): self._plugins: dict[str,WorkerPlugin]={} + + def register(self,plugin:WorkerPlugin) -> None: + if not plugin.name or plugin.name in self._plugins: raise ValueError(f"duplicate or empty worker plugin {plugin.name!r}") + if not plugin.module.startswith("extra.radeon_forge.backends."): + raise ValueError("worker plugins must resolve to an in-tree local Radeon Forge backend module") + self._plugins[plugin.name]=plugin + + def get(self,name:str) -> WorkerPlugin: + try:return self._plugins[name] + except KeyError as exc:raise KeyError(f"unknown local worker plugin {name!r}; available={sorted(self._plugins)}") from exc + + def list(self) -> list[dict[str,Any]]: + return [{"name":item.name,"module":item.module,"description":item.description, + "required":list(item.required),"optional":list(item.optional),"defaults":dict(item.defaults)} + for item in sorted(self._plugins.values(),key=lambda x:x.name)] + + +def default_worker_plugins() -> WorkerPluginRegistry: + registry=WorkerPluginRegistry() + registry.register(WorkerPlugin( + name="legacy-llama", + module="extra.radeon_forge.backends.tinygrad_llama_worker", + description="Resident tinygrad Llama-family worker with exact KV accounting, stage hooks and kernel profiling.", + required=("model",), + optional=("tokenizer","size","quantize","max_context","seed"), + defaults={"size":"1B","max_context":8192,"seed":42}, + )) + return registry diff --git a/extra/radeon_forge/backends/stage_hooks.py b/extra/radeon_forge/backends/stage_hooks.py new file mode 100644 index 0000000000000..1cc9f1c3929ce --- /dev/null +++ b/extra/radeon_forge/backends/stage_hooks.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import hashlib, importlib.util, json +from dataclasses import asdict +from pathlib import Path +from typing import Any, Mapping, Sequence + +from tinygrad import TinyJit + +from ..synthesis.hooks import ExecutionContext, StagePredicate + + +class StageHookError(RuntimeError): pass + + +def _matching_hooks(hooks: Sequence[Mapping[str, Any]], context: ExecutionContext) -> list[Mapping[str, Any]]: + """Resolve the highest-priority matching hook in each exclusive group.""" + selected: dict[str, Mapping[str, Any]] = {} + for hook in hooks: + descriptor = hook.get("descriptor", {}) + if not isinstance(descriptor, Mapping): continue + try: predicate = StagePredicate.from_mapping(descriptor.get("when", {})) + except Exception: continue + if not predicate.matches(context): continue + group = str(descriptor.get("exclusive_group", f"{descriptor.get('layer')}:{descriptor.get('target')}")) + previous = selected.get(group) + if previous is None or int(descriptor.get("priority", 0)) > int(previous.get("descriptor", {}).get("priority", 0)): + selected[group] = hook + return sorted(selected.values(), key=lambda x: -int(x.get("descriptor", {}).get("priority", 0))) + + +def _signature(hook: Mapping[str, Any]) -> str: + parameters = hook.get("parameters", {}) + digest = hashlib.sha256(json.dumps(parameters, sort_keys=True, default=str).encode()).hexdigest()[:10] + return f"{hook.get('activation_id')}:{digest}" + + +class ModelStageHookRuntime: + """Apply validated Python model hooks only in their declared execution states. + + Candidate code runs inside the already-isolated local model subprocess. The + original model objects are retained and restored on every transition or + failure. This adapter never treats activation metadata as validation evidence. + """ + def __init__(self, model: Any): + self.model = model + self._original_layers = tuple(getattr(model, "layers", ())) + self._active_signature: tuple[str, ...] = () + self._loaded_modules: dict[tuple[str, str], Any] = {} + + def _reset(self) -> None: + if self._original_layers and hasattr(self.model, "layers"): self.model.layers = list(self._original_layers) + if hasattr(self.model, "forward_jit") and self.model.forward_jit is not None: self.model.forward_jit = TinyJit(self.model.forward) + self._active_signature = () + + @staticmethod + def _indices(selector: Mapping[str, Any], count: int) -> list[int]: + value = selector.get("indices", "all") + if value == "all": return list(range(count)) + if isinstance(value, int): values = [value] + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): values = [int(x) for x in value] + else: raise StageHookError("transformer block selector indices must be 'all', an integer, or an array") + if any(index < 0 or index >= count for index in values): raise StageHookError("transformer block selector is out of range") + return values + + def _module(self, hook: Mapping[str, Any]) -> Any: + path = Path(str(hook["source_path"])) + expected = str(hook["source_sha256"]) + if not path.is_file(): raise StageHookError(f"hook source is missing: {path}") + actual = hashlib.sha256(path.read_bytes()).hexdigest() + if actual != expected: raise StageHookError("hook source changed after activation") + key = (str(path), expected) + if key in self._loaded_modules: return self._loaded_modules[key] + spec = importlib.util.spec_from_file_location(f"radeon_forge_hook_{expected[:16]}", path) + if spec is None or spec.loader is None: raise StageHookError(f"cannot import hook source {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + self._loaded_modules[key] = module + return module + + def apply(self, hooks: Sequence[Mapping[str, Any]], context: ExecutionContext) -> dict[str, Any]: + selected = [hook for hook in _matching_hooks(hooks, context) + if hook.get("descriptor", {}).get("adapter") == "python_transformer_block"] + signature = tuple(_signature(hook) for hook in selected) + if signature == self._active_signature: + return {"changed": False, "active": [str(hook.get("activation_id")) for hook in selected], "stage": context.stage.value} + self._reset() + if not selected: return {"changed": True, "active": [], "stage": context.stage.value, "baseline": True} + try: + layers = list(self._original_layers) + selected_parameters: dict[str, Mapping[str, Any]] = {} + for hook in selected: + descriptor = hook["descriptor"] + if descriptor.get("mode", "replace") != "replace": raise StageHookError("python_transformer_block currently supports replace mode only") + module = self._module(hook) + parameters = hook.get("parameters", {}) + if not isinstance(parameters, Mapping): raise StageHookError("hook parameters must be a table") + parameters = dict(parameters) + setattr(module, "RADEON_FORGE_PARAMETERS", parameters) + configure = getattr(module, "configure", None) + if configure is not None: + if not callable(configure): raise StageHookError("optional configure attribute must be callable") + configure(parameters) + factory = getattr(module, "build_replacement", None) + if not callable(factory): raise StageHookError("hook must define build_replacement(original, layer_index, model, context)") + selector = descriptor.get("selector", {}) + for index in self._indices(selector if isinstance(selector, Mapping) else {}, len(layers)): + replacement = factory(layers[index], index, self.model, asdict(context)) + if not callable(replacement): raise StageHookError(f"replacement for layer {index} is not callable") + layers[index] = replacement + selected_parameters[str(hook.get("activation_id"))] = parameters + self.model.layers = layers + if hasattr(self.model, "forward_jit") and self.model.forward_jit is not None: self.model.forward_jit = TinyJit(self.model.forward) + self._active_signature = signature + return {"changed": True, "active": [str(hook.get("activation_id")) for hook in selected], + "parameters": selected_parameters, "stage": context.stage.value, "baseline": False} + except Exception as exc: + self._reset() + return {"changed": True, "active": [], "stage": context.stage.value, "baseline": True, + "rolled_back": True, "error": str(exc), "error_type": type(exc).__name__} + + def close(self) -> None: self._reset() diff --git a/extra/radeon_forge/backends/tinygrad_llama_worker.py b/extra/radeon_forge/backends/tinygrad_llama_worker.py new file mode 100644 index 0000000000000..8a105e20ea15c --- /dev/null +++ b/extra/radeon_forge/backends/tinygrad_llama_worker.py @@ -0,0 +1,279 @@ +from __future__ import annotations + +import argparse, hashlib, json, subprocess, sys, time +from pathlib import Path +from typing import Any, Callable, Mapping + +from tinygrad import Context, Device, GlobalCounters, Tensor +from tinygrad.device import Compiled +from tinygrad.nn.state import get_parameters +from examples.llama3 import Tokenizer, build_transformer + +from ..runtime.tool_stream import ToolStreamParser +from ..synthesis.hooks import ExecutionContext, ExecutionStage +from .kv_state import KVReuseLedger +from .stage_hooks import ModelStageHookRuntime + + +MAX_PROFILE_EVENTS_PER_PHASE = 8192 + + +def send(payload): + sys.stdout.write(json.dumps(payload, default=str) + "\n") + sys.stdout.flush() + + +def _name(value: Any) -> str: return str(getattr(value, "display_name", value)) + + +def _collect_profile_events(start: int) -> list[dict[str, Any]]: + """Flatten tinygrad device profile ranges into portable kernel evidence.""" + raw = list(Compiled.profile_events[start:]) + del Compiled.profile_events[start:] + ret: list[dict[str, Any]] = [] + for event in raw: + if hasattr(event, "ents") and hasattr(event, "sigs"): + for entry in event.ents: + try: duration_us = float(event.sigs[entry.en_id] - event.sigs[entry.st_id]) + except Exception: continue + device = str(getattr(entry, "device", "UNKNOWN")) + if device in {"CPU", "TINY"} or duration_us < 0: continue + ret.append({"name": _name(getattr(entry, "name", "kernel")), "device": device, + "duration_ms": duration_us / 1000.0, "profile_event_type": type(event).__name__, + "source": "tinygrad.Compiled.profile_events"}) + elif hasattr(event, "st") and hasattr(event, "en") and getattr(event, "en") is not None: + try: duration_us = float(event.en - event.st) + except Exception: continue + device = str(getattr(event, "device", "UNKNOWN")) + if device in {"CPU", "TINY"} or duration_us < 0: continue + ret.append({"name": _name(getattr(event, "name", "kernel")), "device": device, + "duration_ms": duration_us / 1000.0, "profile_event_type": type(event).__name__, + "source": "tinygrad.Compiled.profile_events"}) + return ret + + +def _profiled(fn: Callable[[], Any]) -> tuple[Any, list[dict[str, Any]]]: + start = len(Compiled.profile_events) + try: + with Context(PROFILE=1): result = fn() + except BaseException: + _collect_profile_events(start) + raise + return result, _collect_profile_events(start) + + +def _emit_kernel_events(events: list[dict[str, Any]], *, stage: str, token_index: int | None = None) -> tuple[int, bool]: + truncated = len(events) > MAX_PROFILE_EVENTS_PER_PHASE + emitted = events[:MAX_PROFILE_EVENTS_PER_PHASE] + for sequence, event in enumerate(emitted): + metrics = {**event, "stage": stage, "sequence": sequence} + if token_index is not None: metrics["token_index"] = token_index + send({"kind": "kernel", "metrics": metrics}) + return len(emitted), truncated + + +def _weight_identity(path: Path, size: str, quantize: str | None) -> str: + try: + stat = path.stat() + identity = f"{path.resolve()}:{stat.st_size}:{stat.st_mtime_ns}:{size}:{quantize}" + except OSError: identity = f"{path}:{size}:{quantize}" + return hashlib.sha256(identity.encode()).hexdigest() + + +def _shape(value: Any) -> list[int | str]: + try: return [int(item) for item in value.shape] + except Exception: return [str(getattr(value, "shape", "unknown"))] + + +def _architecture_identity(model: Any, size: str, quantize: str | None) -> tuple[dict[str, Any], str]: + first = model.layers[0] if getattr(model, "layers", None) else None + attention = getattr(first, "attention", None) + feed_forward = getattr(first, "feed_forward", None) + payload = { + "family": "llama", "size_label": size, "quantize": quantize or "model_default", + "layers": len(getattr(model, "layers", ())), "max_context": int(getattr(model, "max_context", 0)), + "n_heads": int(getattr(attention, "n_heads", 0)), "n_kv_heads": int(getattr(attention, "n_kv_heads", 0)), + "head_dim": int(getattr(attention, "head_dim", 0)), + "embedding_shape": _shape(getattr(getattr(model, "tok_embeddings", None), "weight", None)), + "output_shape": _shape(getattr(getattr(model, "output", None), "weight", None)), + "ffn_w1_shape": _shape(getattr(getattr(feed_forward, "w1", None), "weight", None)), + } + return payload, hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest() + + +def _git_revision() -> str: + try: + root = Path(__file__).resolve().parents[3] + return subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=root, text=True, stderr=subprocess.DEVNULL).strip() + except Exception: return "unknown" + + +def _hook_event(runtime: ModelStageHookRuntime, hooks: list[Mapping[str, Any]], context: ExecutionContext) -> None: + result = runtime.apply(hooks, context) + if result.get("changed") or result.get("rolled_back"): + send({"kind": "hook", "metrics": {**result, "context": context.flattened()}}) + + +def _prefill_without_output_head(model: Any, token_ids: list[int], start_pos: int, device: str) -> None: + """Populate KV state for a prompt chunk without projecting every prompt token to vocabulary logits.""" + if not token_ids: return + tokens = Tensor([token_ids], device=device) + _, seqlen = tokens.shape + h = model.tok_embeddings(tokens).contiguous() + freqs_cis = model.freqs_cis.cast(h.dtype)[:, start_pos:start_pos+seqlen, :, :, :] + mask = (Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1) + if model.max_context != 0 and seqlen > 1 else None) + for layer in model.layers: h = layer(h, start_pos, freqs_cis, mask) + h.realize() + + +def _available_tool_names(tools: Any) -> list[str]: + names = [] + if not isinstance(tools, list): return names + for item in tools: + if not isinstance(item, Mapping): continue + function = item.get("function", {}) + if isinstance(function, Mapping) and isinstance(function.get("name"), str): names.append(str(function["name"])) + return names + + +def main() -> None: + parser = argparse.ArgumentParser(description="Persistent tinygrad Llama backend for Radeon Forge") + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path) + parser.add_argument("--size", choices=("1B", "8B", "70B", "405B"), default="1B") + parser.add_argument("--quantize", choices=("int8", "nf4", "float16", "fp8")) + parser.add_argument("--max-context", type=int, default=8192) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + if not args.model.exists(): raise FileNotFoundError(args.model) + tokenizer_path = args.tokenizer or ((args.model if args.model.is_dir() else args.model.parent) / "tokenizer.model") + if not tokenizer_path.exists(): raise FileNotFoundError(tokenizer_path) + Tensor.manual_seed(args.seed) + + tokenizer = Tokenizer(str(tokenizer_path)) + device = Device.DEFAULT + model = build_transformer(args.model, model_size=args.size, quantize=args.quantize, device=device, max_context=args.max_context) + param_bytes = sum(x.nbytes() for x in get_parameters(model)) + architecture_data, architecture_hash = _architecture_identity(model, args.size, args.quantize) + hook_runtime = ModelStageHookRuntime(model) + kv = KVReuseLedger() + device_obj = Device[device] + architecture = str(getattr(device_obj, "arch", "")) + + def encode_role(role: str) -> list[int]: + return [tokenizer.special_tokens["<|start_header_id|>"]] + tokenizer.encode(role) + [tokenizer.special_tokens["<|end_header_id|>"]] + tokenizer.encode("\n\n") + + def encode_message(message) -> list[int]: + content = message.get("content", "") + if not isinstance(content, str): content = json.dumps(content, default=str) + return encode_role(str(message.get("role", "user"))) + tokenizer.encode(content.strip()) + [tokenizer.special_tokens["<|eot_id|>"]] + + send({"kind": "ready", "name": f"tinygrad-llama-{args.size}", "device": str(device), "gpu": str(device), + "architecture": architecture, "runtime": "tinygrad", "runtime_revision": _git_revision(), + "model": str(args.model), "model_family": "llama", "model_hash": _weight_identity(args.model, args.size, args.quantize), + "model_architecture_hash": architecture_hash, "model_architecture": architecture_data, + "dtype": args.quantize or "model_default", "shapes": {"batch": 1, "max_context": args.max_context}, + "parameter_bytes": param_bytes, "capabilities": {"streaming": True, "structured_tools": True, + "prefix_cache": True, "persistent_kv": True, "kernel_metrics": True, "cancellation": False, + "batched_prefill": True, "stage_hooks": True, "local_only": True}}) + + for line in sys.stdin: + try: + payload = json.loads(line) + if payload.get("op") == "shutdown": + hook_runtime.close() + kv.reset() + return + if payload.get("op") != "generate": raise ValueError("unsupported operation") + request = payload["request"] + metadata = request.get("metadata", {}) + hooks = metadata.get("active_hooks", []) + if not isinstance(hooks, list): hooks = [] + session_id = str(request["session_id"]) + messages = request["messages"] + max_tokens = max(1, min(int(request.get("max_tokens", 256)), args.max_context)) + temperature = float(request.get("temperature", 0.0)) + tool_round = int(metadata.get("tool_round", 0)) + resume_after_tool = bool(metadata.get("resume_after_tool", False)) + tool_names = _available_tool_names(request.get("tools", [])) + tool_parser = ToolStreamParser(tool_names) if tool_names else None + + prompt = [tokenizer.bos_id] + for message in messages: prompt += encode_message(message) + if not messages or messages[-1].get("role") != "assistant": prompt += encode_role("assistant") + if len(prompt) >= args.max_context: raise ValueError(f"prompt has {len(prompt)} tokens, max context is {args.max_context}") + + common = kv.begin(session_id, prompt) + prefill_context = ExecutionContext(ExecutionStage.PREFILL, batch_size=1, prompt_tokens=len(prompt), + context_tokens=max(0, len(prompt)-1), generated_token_index=-1, prefix_reused_tokens=common, + tool_round=tool_round, warm=bool(kv.cached_tokens), attributes={"resume_after_tool": resume_after_tool, "session_id": session_id}) + _hook_event(hook_runtime, hooks, prefill_context) + + prefill_ids = list(prompt[common:-1]) + GlobalCounters.reset() + prefill_start = time.perf_counter_ns() + _, prefill_profile = _profiled(lambda: _prefill_without_output_head(model, prefill_ids, common, device)) + prefill_wall_ms = (time.perf_counter_ns() - prefill_start) / 1e6 + prefill_gpu_ms = GlobalCounters.time_sum_s * 1e3 + kv.commit_prefill(prompt[:-1]) + send({"kind": "prefill", "metrics": {"wall_ms": prefill_wall_ms, "gpu_ms": prefill_gpu_ms, + "prompt_tokens": len(prompt), "prefix_reused_tokens": common, "new_prompt_tokens": len(prefill_ids), + "prefill_chunk_tokens": len(prefill_ids), "prefix_cache_hit_ratio": common / max(1, len(prompt)-1), + "kernel_count": GlobalCounters.kernel_count, "global_mem_bytes": GlobalCounters.global_mem, + "global_ops": GlobalCounters.global_ops, "profile_kernel_events": len(prefill_profile), + "resume_after_tool": resume_after_tool, "output_head_skipped": True}}) + emitted, prefill_truncated = _emit_kernel_events(prefill_profile, stage="prefill") + if prefill_truncated: + send({"kind": "metric", "metrics": {"name": "profile_truncated", "stage": "prefill", + "captured": emitted, "available": len(prefill_profile)}}) + + start_pos, last_tok = len(prompt) - 1, prompt[-1] + generated = 0 + for index in range(max_tokens): + stage = ExecutionStage.FIRST_TOKEN if index == 0 else ExecutionStage.DECODE + decode_context = ExecutionContext(stage, batch_size=1, prompt_tokens=len(prompt), context_tokens=start_pos, + generated_token_index=index, prefix_reused_tokens=common, tool_round=tool_round, warm=True, + attributes={"resume_after_tool": resume_after_tool, "session_id": session_id}) + _hook_event(hook_runtime, hooks, decode_context) + GlobalCounters.reset() + wall_start = time.perf_counter_ns() + input_tok = last_tok + tok, profile = _profiled(lambda: model(Tensor([[input_tok]], device=device), start_pos, temperature, 0, 0.0, 0.0, 0.0).item()) + wall_ms = (time.perf_counter_ns() - wall_start) / 1e6 + gpu_ms = GlobalCounters.time_sum_s * 1e3 + kv.commit_decode_input(input_tok, start_pos) + start_pos += 1 + last_tok = tok + if tok in tokenizer.stop_tokens: + if tool_parser is not None: tool_parser.finalize() + send({"kind": "done", "finish_reason": "stop", "metrics": {"generated_tokens": generated, + "materialized_kv_tokens": len(kv.cached_tokens)}}) + break + generated += 1 + piece = tokenizer.decode([tok]) + parsed_tool = tool_parser.feed(piece) if tool_parser is not None else None + send({"kind": "token", "text": piece, "metrics": {"index": index, "stage": stage.value, + "wall_ms": wall_ms, "gpu_ms": gpu_ms, "kernel_count": GlobalCounters.kernel_count, + "global_mem_bytes": GlobalCounters.global_mem, "global_ops": GlobalCounters.global_ops, + "profile_kernel_events": len(profile), "context_tokens": start_pos, + "parameter_bandwidth_gbs": (param_bytes / max(GlobalCounters.time_sum_s, 1e-12)) / 1e9}}) + emitted, truncated = _emit_kernel_events(profile, stage=stage.value, token_index=index) + if truncated: + send({"kind": "metric", "metrics": {"name": "profile_truncated", "stage": stage.value, "token_index": index, + "captured": emitted, "available": len(profile)}}) + if parsed_tool is not None: + send({"kind": "tool_call", "tool_call": parsed_tool.to_event()}) + send({"kind": "done", "finish_reason": "tool_call", "metrics": {"generated_tokens": generated, + "materialized_kv_tokens": len(kv.cached_tokens), "tool_name": parsed_tool.name}}) + break + else: + if tool_parser is not None: tool_parser.finalize() + send({"kind": "done", "finish_reason": "length", "metrics": {"generated_tokens": generated, + "materialized_kv_tokens": len(kv.cached_tokens), "pending_uncached_output_token": True}}) + except Exception as exc: + send({"kind": "error", "error": str(exc), "error_type": type(exc).__name__}) + + +if __name__ == "__main__": main() diff --git a/extra/radeon_forge/cli.py b/extra/radeon_forge/cli.py new file mode 100644 index 0000000000000..c7a1a7db262fd --- /dev/null +++ b/extra/radeon_forge/cli.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +from .command_backend import CommandHarness +from .contracts import Objective, WorkloadContract +from .families import rdna3_asm_matmul_family +from .ledger import ExperimentLedger +from .permissions import Action, PermissionController +from .tuner import successive_halving + + +def _parse_budgets(value: str) -> tuple[int, ...]: + budgets = tuple(int(part) for part in value.split(",") if part.strip()) + if not budgets or any(budget <= 0 for budget in budgets): raise argparse.ArgumentTypeError("budgets must be positive comma-separated integers") + return budgets + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Oracle-guided Radeon Forge autotuner") + parser.add_argument("--root", default=".", help="tinygrad repository root") + parser.add_argument("--family", choices=("rdna3-asm-matmul",), default="rdna3-asm-matmul") + parser.add_argument("--n", type=int, default=1024, help="square GEMM dimension for the first gfx1100 hardware family") + parser.add_argument("--budgets", type=_parse_budgets, default=(3, 10, 30), help="successive-halving trial counts, e.g. 3,10,30") + parser.add_argument("--reduction", type=int, default=4) + parser.add_argument("--ledger", default="evidence/radeon_forge/events.jsonl") + parser.add_argument("--dry-run", action="store_true", help="print the plan without compiling or benchmarking") + parser.add_argument("--approve-benchmark", action="store_true", help="explicitly approve local W7900 benchmark execution") + return parser + + +def main(argv: list[str] | None = None) -> int: + args = build_parser().parse_args(argv) + root = Path(args.root).resolve() + family = rdna3_asm_matmul_family(root, args.n) + candidates = family.candidates() + plan = { + "family": family.name, + "candidate_count": len(candidates), + "budgets": args.budgets, + "reduction": args.reduction, + "objective": Objective.P95_LATENCY_US.value, + "target": "gfx1100", + "reference": "tinygrad matmul", + "candidate_ids": [candidate.candidate_id for candidate in candidates], + } + print(json.dumps({"plan": plan}, indent=2, sort_keys=True)) + if args.dry_run: return 0 + if not args.approve_benchmark: + print("Refusing hardware execution: pass --approve-benchmark after reviewing the plan.", file=sys.stderr) + return 2 + + ledger = ExperimentLedger(root / args.ledger) + ledger.append("optimization_requested", plan) + controller = PermissionController() + # Upper bound: every candidate in every round. The actual tuner consumes fewer + # uses after successive halving removes candidates. + grant = controller.issue([Action.BENCHMARK], f"approved benchmark plan for {family.name}", max_uses=len(candidates) * len(args.budgets)) + harness = CommandHarness([sys.executable, "-m", "extra.radeon_forge.workloads.rdna3_asm_matmul"], cwd=root) + + def approved_evaluator(candidate, budget): + authorized = controller.authorize(grant.token, Action.BENCHMARK) + ledger.append("permission_used", {"action": Action.BENCHMARK.value, "reason": authorized.reason, + "candidate_id": candidate.candidate_id, "budget": budget}) + return harness(candidate, budget) + + contract = WorkloadContract( + name="rdna3-asm-matmul", + target="gfx1100", + objective=Objective.P95_LATENCY_US, + max_abs_error=0.1, + max_rel_error=1000.0, + max_vgprs=192, + max_lds_bytes=32768, + forbid_spills=True, + metadata={"final_submission_metric": False, "purpose": "validate the Forge hardware tuning loop"}, + ) + ledger.append("workload_contract", contract) + summary = successive_halving(candidates, approved_evaluator, contract, budgets=args.budgets, reduction=args.reduction, ledger=ledger) + if summary.winner is None: + print(json.dumps({"status": "no_feasible_candidate", "evaluations": len(summary.evaluated)}, indent=2)) + return 1 + print(json.dumps({"status": "winner", "rounds": summary.rounds, "evaluations": len(summary.evaluated), + "result": summary.winner.to_dict()}, indent=2, sort_keys=True)) + return 0 + + +if __name__ == "__main__": raise SystemExit(main()) diff --git a/extra/radeon_forge/command_backend.py b/extra/radeon_forge/command_backend.py new file mode 100644 index 0000000000000..84bf002b913fc --- /dev/null +++ b/extra/radeon_forge/command_backend.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +import json +import os +import statistics +import subprocess +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, Sequence + +from .contracts import Candidate, CorrectnessReport, ResourceReport, TrialResult + + +def _percentile(values: Sequence[float], q: float) -> float | None: + if not values: return None + ordered = sorted(values) + index = min(len(ordered) - 1, max(0, int((len(ordered) - 1) * q + 0.999999))) + return ordered[index] + + +@dataclass(frozen=True) +class CommandHarness: + """Evaluate candidates through a workload-specific executable. + + The executable receives the candidate and budget through environment + variables and must print one JSON object on its final non-empty stdout line. + This keeps Forge independent of any one kernel launch abstraction while the + reference implementation and benchmark protocol remain explicit. + """ + + command: Sequence[str] + cwd: str | Path | None = None + env: Mapping[str, str] | None = None + timeout_seconds: int = 300 + + def __call__(self, candidate: Candidate, budget: int) -> TrialResult: + run_env = dict(os.environ) + if self.env is not None: run_env.update(self.env) + run_env["RADEON_FORGE_CANDIDATE_JSON"] = json.dumps(candidate.to_dict(), sort_keys=True) + run_env["RADEON_FORGE_BUDGET"] = str(budget) + proc = subprocess.run(list(self.command), cwd=self.cwd, env=run_env, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, + timeout=self.timeout_seconds, check=False) + lines = [line for line in proc.stdout.splitlines() if line.strip()] + if proc.returncode != 0 or not lines: + reason = f"command failed with exit {proc.returncode}: {proc.stderr[-2000:]}" + return TrialResult(candidate, CorrectnessReport(False, reason=reason), compile_ok=False, stable=False, + rejected_reason="command_failed", evidence={"stdout": proc.stdout[-2000:], "stderr": proc.stderr[-2000:]}) + try: + payload: dict[str, Any] = json.loads(lines[-1]) + except json.JSONDecodeError as exc: + return TrialResult(candidate, CorrectnessReport(False, reason=f"invalid JSON result: {exc}"), compile_ok=False, stable=False, + rejected_reason="invalid_result", evidence={"stdout": proc.stdout[-4000:], "stderr": proc.stderr[-2000:]}) + + correctness_data = payload.get("correctness", {}) + resource_data = payload.get("resources", {}) + samples = tuple(float(value) for value in payload.get("samples_us", ())) + correctness = CorrectnessReport( + passed=bool(correctness_data.get("passed", False)), + max_abs_error=float(correctness_data.get("max_abs_error", 0.0)), + max_rel_error=float(correctness_data.get("max_rel_error", 0.0)), + checked_values=int(correctness_data.get("checked_values", 0)), + reason=str(correctness_data.get("reason", "")), + ) + resources = ResourceReport( + vgprs=resource_data.get("vgprs"), sgprs=resource_data.get("sgprs"), lds_bytes=resource_data.get("lds_bytes"), + scratch_bytes=resource_data.get("scratch_bytes"), occupancy=resource_data.get("occupancy"), + spilled_vgprs=int(resource_data.get("spilled_vgprs", 0)), spilled_sgprs=int(resource_data.get("spilled_sgprs", 0)), + ) + return TrialResult( + candidate=candidate, + correctness=correctness, + samples_us=samples, + median_latency_us=float(payload["median_latency_us"]) if "median_latency_us" in payload else (statistics.median(samples) if samples else None), + p95_latency_us=float(payload["p95_latency_us"]) if "p95_latency_us" in payload else _percentile(samples, 0.95), + end_to_end_p95_ms=float(payload["end_to_end_p95_ms"]) if "end_to_end_p95_ms" in payload else None, + resources=resources, + compile_ok=bool(payload.get("compile_ok", True)), + stable=bool(payload.get("stable", True)), + rejected_reason=str(payload.get("rejected_reason", "")), + evidence={"stdout": proc.stdout, "stderr": proc.stderr, **payload.get("evidence", {})}, + ) diff --git a/extra/radeon_forge/compile_backend.py b/extra/radeon_forge/compile_backend.py new file mode 100644 index 0000000000000..f9b6773433c89 --- /dev/null +++ b/extra/radeon_forge/compile_backend.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import shutil +import subprocess +import tempfile +from dataclasses import dataclass +from pathlib import Path + +from tinygrad.runtime.support.compiler_amd import HIPCCCompiler + +from .amd_metadata import resource_report_from_text +from .contracts import Candidate, ResourceReport +from .families import compiler_defines + + +@dataclass(frozen=True) +class CompiledCandidate: + candidate: Candidate + library: bytes + resources: ResourceReport + metadata_text: str + + +def _read_metadata(library: bytes) -> str: + tools_and_args = ( + ("/opt/rocm/llvm/bin/llvm-readobj", "--notes"), + ("/opt/rocm/llvm/bin/llvm-objdump", "--notes"), + ("llvm-readobj", "--notes"), + ("llvm-objdump", "--notes"), + ) + with tempfile.NamedTemporaryFile(suffix=".hsaco") as f: + f.write(library) + f.flush() + outputs: list[str] = [] + for tool, flag in tools_and_args: + executable = tool if Path(tool).exists() else shutil.which(tool) + if executable is None: continue + proc = subprocess.run([executable, flag, f.name], text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, check=False) + if proc.stdout: outputs.append(proc.stdout) + if proc.returncode == 0 and proc.stdout: break + return "\n".join(outputs) + + +def compile_candidate(candidate: Candidate, target: str = "gfx1100") -> CompiledCandidate: + if candidate.source_path is None: raise ValueError(f"candidate {candidate.candidate_id} has no source_path") + source = Path(candidate.source_path).read_text(encoding="utf-8") + options = ["-std=c++20", "-ffast-math", *compiler_defines(candidate.parameters)] + library = HIPCCCompiler(target, options).compile_cached(source) + metadata_text = _read_metadata(library) + return CompiledCandidate(candidate, library, resource_report_from_text(metadata_text), metadata_text) diff --git a/extra/radeon_forge/contracts.py b/extra/radeon_forge/contracts.py new file mode 100644 index 0000000000000..b5e9159b7bafd --- /dev/null +++ b/extra/radeon_forge/contracts.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from enum import Enum +from typing import Any, Mapping + + +class Objective(str, Enum): + MEDIAN_LATENCY_US = "median_latency_us" + P95_LATENCY_US = "p95_latency_us" + END_TO_END_P95_MS = "end_to_end_p95_ms" + + +@dataclass(frozen=True) +class WorkloadContract: + name: str + target: str = "gfx1100" + objective: Objective = Objective.P95_LATENCY_US + max_abs_error: float = 1e-3 + max_rel_error: float = 1e-3 + max_vram_bytes: int | None = None + max_vgprs: int | None = None + max_lds_bytes: int | None = None + forbid_spills: bool = True + metadata: Mapping[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class Candidate: + candidate_id: str + family: str + parameters: Mapping[str, int | float | str | bool] + source_path: str | None = None + hypothesis: str = "" + parent_id: str | None = None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class ResourceReport: + vgprs: int | None = None + sgprs: int | None = None + lds_bytes: int | None = None + scratch_bytes: int | None = None + occupancy: int | None = None + spilled_vgprs: int = 0 + spilled_sgprs: int = 0 + + @property + def has_spills(self) -> bool: + return self.spilled_vgprs > 0 or self.spilled_sgprs > 0 + + +@dataclass(frozen=True) +class CorrectnessReport: + passed: bool + max_abs_error: float = 0.0 + max_rel_error: float = 0.0 + checked_values: int = 0 + reason: str = "" + + +@dataclass(frozen=True) +class TrialResult: + candidate: Candidate + correctness: CorrectnessReport + samples_us: tuple[float, ...] = () + median_latency_us: float | None = None + p95_latency_us: float | None = None + end_to_end_p95_ms: float | None = None + resources: ResourceReport = field(default_factory=ResourceReport) + compile_ok: bool = True + stable: bool = True + rejected_reason: str = "" + evidence: Mapping[str, Any] = field(default_factory=dict) + + def metric(self, objective: Objective) -> float: + value = { + Objective.MEDIAN_LATENCY_US: self.median_latency_us, + Objective.P95_LATENCY_US: self.p95_latency_us, + Objective.END_TO_END_P95_MS: self.end_to_end_p95_ms, + }[objective] + return float("inf") if value is None else value + + def is_feasible(self, contract: WorkloadContract) -> bool: + if not self.compile_ok or not self.stable or not self.correctness.passed: return False + if contract.forbid_spills and self.resources.has_spills: return False + if contract.max_vgprs is not None and self.resources.vgprs is not None and self.resources.vgprs > contract.max_vgprs: return False + if contract.max_lds_bytes is not None and self.resources.lds_bytes is not None and self.resources.lds_bytes > contract.max_lds_bytes: return False + if self.correctness.max_abs_error > contract.max_abs_error: return False + if self.correctness.max_rel_error > contract.max_rel_error: return False + return self.metric(contract.objective) != float("inf") + + def to_dict(self) -> dict[str, Any]: + return asdict(self) diff --git a/extra/radeon_forge/evaluation/__init__.py b/extra/radeon_forge/evaluation/__init__.py new file mode 100644 index 0000000000000..1cdec57d7248a --- /dev/null +++ b/extra/radeon_forge/evaluation/__init__.py @@ -0,0 +1,3 @@ +from .suite import AgentTask, AgentTaskSuite, SuiteResult, TaskResult, run_suite + +__all__=["AgentTask","AgentTaskSuite","SuiteResult","TaskResult","run_suite"] diff --git a/extra/radeon_forge/evaluation/cli.py b/extra/radeon_forge/evaluation/cli.py new file mode 100644 index 0000000000000..9c8eebc24c7d2 --- /dev/null +++ b/extra/radeon_forge/evaluation/cli.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +import argparse, json, sys +from pathlib import Path + +from ..backends.plugins import default_worker_plugins +from ..runtime import ForgeEngine, JsonlProcessBackend +from .suite import AgentTaskSuite, run_suite + + +def main() -> None: + parser=argparse.ArgumentParser(description="Run a frozen private-agent suite on the local Radeon Forge engine") + parser.add_argument("--suite",type=Path,required=True) + parser.add_argument("--workspace",type=Path,default=Path.cwd()) + parser.add_argument("--model-plugin",default="legacy-llama") + parser.add_argument("--model",type=Path,required=True) + parser.add_argument("--tokenizer",type=Path) + parser.add_argument("--size",default="1B") + parser.add_argument("--quantize") + parser.add_argument("--max-context",type=int,default=8192) + parser.add_argument("--seed",type=int,default=42) + parser.add_argument("--output",type=Path) + args=parser.parse_args() + + plugin=default_worker_plugins().get(args.model_plugin) + command=plugin.command({"model":args.model,"tokenizer":args.tokenizer,"size":args.size,"quantize":args.quantize, + "max_context":args.max_context,"seed":args.seed}) + engine=ForgeEngine(JsonlProcessBackend(command),args.workspace) + try:result=run_suite(engine,AgentTaskSuite.load(args.suite)).to_dict() + finally:engine.close() + text=json.dumps(result,indent=2,default=str)+"\n" + if args.output: + args.output.parent.mkdir(parents=True,exist_ok=True);args.output.write_text(text,encoding="utf-8") + sys.stdout.write(text) + raise SystemExit(0 if result["passed_tasks"]==result["total_tasks"] else 2) + + +if __name__=="__main__":main() diff --git a/extra/radeon_forge/evaluation/compare.py b/extra/radeon_forge/evaluation/compare.py new file mode 100644 index 0000000000000..c6408d0be7ad2 --- /dev/null +++ b/extra/radeon_forge/evaluation/compare.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass,asdict +from pathlib import Path +from typing import Any,Mapping + + +class IncomparableSuites(ValueError):pass + + +@dataclass(frozen=True) +class MetricDelta: + baseline:float + candidate:float + absolute:float + relative:float|None + + +@dataclass(frozen=True) +class EvaluationGate: + passed:bool + task_success_preserved:bool + tool_validity_preserved:bool + all_baseline_passes_still_pass:bool + reasons:tuple[str,...] + + +def _load(value:str|Path|Mapping[str,Any])->dict[str,Any]: + if isinstance(value,Mapping):return dict(value) + path=Path(value) + payload=json.loads(path.read_text(encoding="utf-8")) + if not isinstance(payload,dict):raise ValueError("evaluation result must be an object") + return payload + + +def _delta(base:float,candidate:float)->dict[str,Any]: + return asdict(MetricDelta(base,candidate,candidate-base,None if base==0 else (candidate-base)/base)) + + +def _tasks(payload:Mapping[str,Any])->dict[str,Mapping[str,Any]]: + rows=payload.get("results",()) + if not isinstance(rows,list):raise ValueError("evaluation results must be an array") + ret={} + for row in rows: + if not isinstance(row,Mapping) or not row.get("task_id"):continue + ret[str(row["task_id"])]=row + return ret + + +def compare_evaluations(baseline_value:str|Path|Mapping[str,Any],candidate_value:str|Path|Mapping[str,Any], + *,max_task_success_drop:float=0.0,max_tool_validity_drop:float=0.0)->dict[str,Any]: + baseline,candidate=_load(baseline_value),_load(candidate_value) + if baseline.get("suite_hash")!=candidate.get("suite_hash"): + raise IncomparableSuites(f"suite_hash differs: baseline={baseline.get('suite_hash')} candidate={candidate.get('suite_hash')}") + base_tasks,cand_tasks=_tasks(baseline),_tasks(candidate) + if set(base_tasks)!=set(cand_tasks): + raise IncomparableSuites(f"task ids differ: baseline={sorted(base_tasks)} candidate={sorted(cand_tasks)}") + + base_success=float(baseline.get("task_success_rate",0.0));cand_success=float(candidate.get("task_success_rate",0.0)) + base_tools=float(baseline.get("tool_call_validity_rate",0.0));cand_tools=float(candidate.get("tool_call_validity_rate",0.0)) + success_ok=cand_success+max_task_success_drop>=base_success + tools_ok=cand_tools+max_tool_validity_drop>=base_tools + regressed=[task_id for task_id,row in base_tasks.items() if bool(row.get("passed")) and not bool(cand_tasks[task_id].get("passed"))] + reasons=[] + if not success_ok:reasons.append(f"task success regressed {base_success:.4f}->{cand_success:.4f}") + if not tools_ok:reasons.append(f"tool validity regressed {base_tools:.4f}->{cand_tools:.4f}") + if regressed:reasons.append(f"baseline-passing tasks regressed: {regressed}") + gate=EvaluationGate(not reasons,success_ok,tools_ok,not regressed,tuple(reasons)) + + per_task=[] + for task_id in sorted(base_tasks): + left,right=base_tasks[task_id],cand_tasks[task_id] + per_task.append({"task_id":task_id,"baseline_passed":bool(left.get("passed")),"candidate_passed":bool(right.get("passed")), + "latency_ms":_delta(float(left.get("latency_ms",0.0)),float(right.get("latency_ms",0.0))), + "baseline_tools":left.get("observed_tools",[]),"candidate_tools":right.get("observed_tools",[]), + "candidate_failure_reasons":right.get("failure_reasons",[])}) + + latency={name:_delta(float(baseline.get(name,0.0)),float(candidate.get(name,0.0))) + for name in ("p50_latency_ms","p95_latency_ms","mean_latency_ms")} + return {"suite_name":baseline.get("suite_name"),"suite_hash":baseline.get("suite_hash"),"gate":asdict(gate), + "quality":{"task_success_rate":_delta(base_success,cand_success),"tool_call_validity_rate":_delta(base_tools,cand_tools)}, + "latency":latency,"per_task":per_task, + "deployable_speedup":gate.passed and latency["p95_latency_ms"]["absolute"]<0, + "interpretation":"Latency improvement is deployable only when the frozen quality/tool gates pass."} diff --git a/extra/radeon_forge/evaluation/suite.py b/extra/radeon_forge/evaluation/suite.py new file mode 100644 index 0000000000000..f026fe2d61f35 --- /dev/null +++ b/extra/radeon_forge/evaluation/suite.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +import hashlib, json, math, re, statistics, time +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any, Mapping, Sequence + +from ..runtime import ForgeEngine +from ..runtime.session import SessionState + + +@dataclass(frozen=True) +class AgentTask: + task_id: str + prompt: str + expected_tools: tuple[str, ...] = () + allowed_tools: tuple[str, ...] = () + required_output_regex: tuple[str, ...] = () + forbidden_output_regex: tuple[str, ...] = () + max_tokens: int = 256 + max_agent_steps: int = 8 + timeout_seconds: float = 120.0 + metadata: Mapping[str, Any] = field(default_factory=dict) + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> AgentTask: + task_id, prompt = str(value.get("task_id", "")).strip(), str(value.get("prompt", "")).strip() + if not task_id or not prompt: raise ValueError("task_id and prompt are required") + def strings(name: str) -> tuple[str, ...]: + raw=value.get(name,()) + if not isinstance(raw,list) or not all(isinstance(item,str) for item in raw): raise ValueError(f"{name} must be a string array") + return tuple(raw) + return cls(task_id,prompt,strings("expected_tools"),strings("allowed_tools"),strings("required_output_regex"), + strings("forbidden_output_regex"),int(value.get("max_tokens",256)),int(value.get("max_agent_steps",8)), + float(value.get("timeout_seconds",120.0)),dict(value.get("metadata",{}))) + + +@dataclass(frozen=True) +class AgentTaskSuite: + name: str + tasks: tuple[AgentTask, ...] + description: str = "" + metadata: Mapping[str, Any] = field(default_factory=dict) + + @property + def suite_hash(self) -> str: + payload=asdict(self) + return hashlib.sha256(json.dumps(payload,sort_keys=True,separators=(",",":"),default=str).encode()).hexdigest() + + @classmethod + def load(cls,path: str|Path) -> AgentTaskSuite: + payload=json.loads(Path(path).read_text(encoding="utf-8")) + if int(payload.get("format_version",0)) != 1: raise ValueError("unsupported agent suite format") + tasks=tuple(AgentTask.from_mapping(item) for item in payload.get("tasks",())) + if not tasks: raise ValueError("suite requires at least one task") + identifiers=[task.task_id for task in tasks] + if len(identifiers)!=len(set(identifiers)): raise ValueError("task ids must be unique") + return cls(str(payload.get("name","")).strip() or Path(path).stem,tasks,str(payload.get("description","")),dict(payload.get("metadata",{}))) + + +@dataclass(frozen=True) +class TaskResult: + task_id: str + passed: bool + latency_ms: float + final_output: str + observed_tools: tuple[str, ...] + expected_tools: tuple[str, ...] + tool_call_valid: bool + output_valid: bool + terminal_state: str + generated_tokens: int + prompt_tokens: int + failure_reasons: tuple[str, ...] + session_id: str + trace_id: str + + +@dataclass(frozen=True) +class SuiteResult: + suite_name: str + suite_hash: str + passed_tasks: int + total_tasks: int + task_success_rate: float + tool_call_validity_rate: float + p50_latency_ms: float + p95_latency_ms: float + mean_latency_ms: float + results: tuple[TaskResult, ...] + + def to_dict(self) -> dict[str, Any]: return asdict(self) + + +def _percentile(values: Sequence[float], fraction: float) -> float: + if not values:return 0.0 + ordered=sorted(values); index=max(0,min(len(ordered)-1,math.ceil(len(ordered)*fraction)-1)) + return ordered[index] + + +def _last_assistant(messages: Sequence[Mapping[str, Any]]) -> str: + for message in reversed(messages): + if message.get("role")=="assistant" and "" not in str(message.get("content","")): + return str(message.get("content","")) + return "" + + +def _tool_sequence(session) -> tuple[str, ...]: + names=[] + for event in session.events: + if event.kind=="tool_started": + call=event.data.get("call",{}) + if isinstance(call,Mapping) and call.get("name"):names.append(str(call["name"])) + return tuple(names) + + +def _token_usage(session) -> tuple[int,int]: + prompt=generated=0 + for event in session.events: + if event.kind=="prefill":prompt=max(prompt,int(event.data.get("prompt_tokens",0))) + elif event.kind=="generation_done":generated+=int((event.data.get("metrics",{}) or {}).get("generated_tokens",0)) + return prompt,generated + + +def run_task(engine: ForgeEngine,task: AgentTask) -> TaskResult: + session=engine.create_session(); session.max_agent_steps=task.max_agent_steps + started=time.perf_counter_ns(); deadline=time.monotonic()+task.timeout_seconds + failures=[] + try: engine.run_message(session.session_id,task.prompt,task.max_tokens,0.0) + except Exception as exc: failures.append(f"initial generation failed: {type(exc).__name__}: {exc}") + while session.state is SessionState.AWAITING_TOOL_APPROVAL and not failures: + if time.monotonic()>=deadline: + failures.append("task timeout while awaiting tool execution");break + call=session.pending_tool_call + if call is None: + failures.append("session requested approval without a pending tool");break + if call.name not in task.allowed_tools: + failures.append(f"tool {call.name!r} is not allowed by the frozen task policy") + session.reject_tool("not allowed by frozen evaluation policy") + break + try:engine.run_tool_approval(session.session_id,f"Frozen suite permits local tool {call.name}",task.max_tokens) + except Exception as exc: + failures.append(f"tool/resume failed: {type(exc).__name__}: {exc}");break + latency_ms=(time.perf_counter_ns()-started)/1e6 + observed=_tool_sequence(session) + tool_valid=observed==task.expected_tools and all(name in task.allowed_tools for name in observed) + if not tool_valid:failures.append(f"tool sequence expected={task.expected_tools!r} observed={observed!r}") + output=_last_assistant(session.messages) + output_valid=True + for pattern in task.required_output_regex: + if re.search(pattern,output,re.I|re.S) is None: + output_valid=False;failures.append(f"required output pattern did not match: {pattern!r}") + for pattern in task.forbidden_output_regex: + if re.search(pattern,output,re.I|re.S) is not None: + output_valid=False;failures.append(f"forbidden output pattern matched: {pattern!r}") + if session.state not in {SessionState.COMPLETED,SessionState.CANCELLED}: + failures.append(f"non-terminal-success session state: {session.state.value}") + if latency_ms>task.timeout_seconds*1000:failures.append("task exceeded latency timeout") + prompt_tokens,generated_tokens=_token_usage(session) + passed=not failures and tool_valid and output_valid and session.state is SessionState.COMPLETED + return TaskResult(task.task_id,passed,latency_ms,output,observed,task.expected_tools,tool_valid,output_valid, + session.state.value,generated_tokens,prompt_tokens,tuple(failures),session.session_id,session.trace.trace_id) + + +def run_suite(engine: ForgeEngine,suite: AgentTaskSuite) -> SuiteResult: + results=tuple(run_task(engine,task) for task in suite.tasks) + latencies=[item.latency_ms for item in results] + return SuiteResult(suite.name,suite.suite_hash,sum(item.passed for item in results),len(results), + sum(item.passed for item in results)/len(results),sum(item.tool_call_valid for item in results)/len(results), + statistics.median(latencies),_percentile(latencies,0.95),statistics.mean(latencies),results) diff --git a/extra/radeon_forge/families.py b/extra/radeon_forge/families.py new file mode 100644 index 0000000000000..6a21d4cf8d999 --- /dev/null +++ b/extra/radeon_forge/families.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import hashlib +import itertools +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, Sequence + +from .contracts import Candidate + + +@dataclass(frozen=True) +class KernelFamily: + name: str + source_path: Path + fixed_defines: Mapping[str, int | float | str | bool] + search_space: Mapping[str, Sequence[int | float | str | bool]] + hypothesis: str + + def candidates(self) -> list[Candidate]: + names = tuple(self.search_space) + values = [self.search_space[name] for name in names] + candidates: list[Candidate] = [] + for combination in itertools.product(*values): + params: dict[str, int | float | str | bool] = dict(self.fixed_defines) + params.update(zip(names, combination)) + digest = hashlib.sha256(repr(sorted(params.items())).encode()).hexdigest()[:10] + candidates.append(Candidate(candidate_id=f"{self.name}-{digest}", family=self.name, parameters=params, + source_path=str(self.source_path), hypothesis=self.hypothesis)) + return candidates + + +def compiler_defines(parameters: Mapping[str, Any]) -> list[str]: + defines: list[str] = [] + for key, value in sorted(parameters.items()): + if isinstance(value, bool): value = int(value) + defines.append(f"-D{key}={value}") + return defines + + +def rdna3_asm_matmul_family(root: str | Path, n: int = 1024) -> KernelFamily: + """Autotuning family around tinygrad's native gfx1100 AMD-DSL GEMM. + + Each FMAC ordering contains exactly the same mathematical operations; only + instruction order changes. LIMIT_OCC controls the deliberate LDS occupancy + limiter already present in the kernel harness. Correctness is still checked + for every candidate against tinygrad matmul. + """ + source = Path(root) / "extra/gemm/amd_asm_matmul.py" + return KernelFamily( + name="rdna3-asm-matmul", + source_path=source, + fixed_defines={"N": n}, + search_space={ + "FMAC_ORDER": ("optimized", "row_major", "column_major", "snake"), + "LIMIT_OCC": (2, 4, 8), + }, + hypothesis="Search VOPD FMAC issue order and occupancy trade-offs for the existing native gfx1100 GEMM.", + ) + + +def rdna3_rmsnorm_fp8_family(root: str | Path, n_elems: int, hidden: int, eps_literal: str = "1e-5f") -> KernelFamily: + """Seed family from tinygrad's fused RMSNorm/multiply/FP8 kernel. + + The source was originally tuned for a different AMD target. Forge treats it + as a pattern and recompiles every candidate for gfx1100; unsupported or + resource-heavy candidates are rejected before benchmarking. + """ + source = Path(root) / "extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/fused_rmsnorm_mul_quantize_fp8.cpp" + return KernelFamily( + name="rdna3-fused-rmsnorm-mul-quantize-fp8", + source_path=source, + fixed_defines={"N_ELEMS": n_elems, "HIDDEN": hidden, "EPS_LITERAL": eps_literal, "HAS_RESIDUAL": 0}, + search_space={"NUM_WG": (128, 256, 512, 1024), "THREADS_PER_WG": (64, 128, 256)}, + hypothesis="Tune work distribution for gfx1100 while preserving the fused single-HBM-pass algorithm.", + ) diff --git a/extra/radeon_forge/knowledge.py b/extra/radeon_forge/knowledge.py new file mode 100644 index 0000000000000..a0ed8af401755 --- /dev/null +++ b/extra/radeon_forge/knowledge.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import math +import re +from collections import Counter +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable, Sequence + + +_TOKEN = re.compile(r"[A-Za-z0-9_+#.-]+") + + +@dataclass(frozen=True) +class KnowledgeChunk: + source: str + start_line: int + end_line: int + text: str + + +@dataclass(frozen=True) +class RetrievalHit: + chunk: KnowledgeChunk + score: float + + @property + def citation(self) -> str: + return f"{self.chunk.source}:{self.chunk.start_line}-{self.chunk.end_line}" + + +class LocalKnowledgeBase: + """Small deterministic local retriever for architecture and experiment knowledge. + + This deliberately has no network path. It is not intended to replace a full + embedding index; its job is to provide auditable local excerpts and citations + to the planning model for the supported project corpus. + """ + + def __init__(self, chunks: Sequence[KnowledgeChunk]): + self.chunks = tuple(chunks) + self._term_counts = [Counter(self._tokens(chunk.text)) for chunk in self.chunks] + document_frequency: Counter[str] = Counter() + for counts in self._term_counts: document_frequency.update(counts.keys()) + self._idf = {term: math.log((1 + len(self.chunks)) / (1 + frequency)) + 1.0 for term, frequency in document_frequency.items()} + + @staticmethod + def _tokens(text: str) -> list[str]: + return [token.lower() for token in _TOKEN.findall(text)] + + @classmethod + def from_paths(cls, paths: Iterable[str | Path], lines_per_chunk: int = 40, overlap: int = 5) -> "LocalKnowledgeBase": + if lines_per_chunk <= 0 or overlap < 0 or overlap >= lines_per_chunk: raise ValueError("invalid chunk dimensions") + chunks: list[KnowledgeChunk] = [] + step = lines_per_chunk - overlap + for path_value in paths: + path = Path(path_value) + lines = path.read_text(encoding="utf-8").splitlines() + for start in range(0, len(lines), step): + selected = lines[start:start + lines_per_chunk] + if not selected: continue + chunks.append(KnowledgeChunk(str(path), start + 1, start + len(selected), "\n".join(selected))) + return cls(chunks) + + def search(self, query: str, top_k: int = 5) -> tuple[RetrievalHit, ...]: + if top_k <= 0: raise ValueError("top_k must be positive") + query_counts = Counter(self._tokens(query)) + if not query_counts: return () + hits: list[RetrievalHit] = [] + for chunk, counts in zip(self.chunks, self._term_counts): + length_norm = max(1.0, math.sqrt(sum(value * value for value in counts.values()))) + score = sum(query_count * counts.get(term, 0) * self._idf.get(term, 1.0) for term, query_count in query_counts.items()) / length_norm + if score > 0: hits.append(RetrievalHit(chunk, score)) + hits.sort(key=lambda hit: (-hit.score, hit.chunk.source, hit.chunk.start_line)) + return tuple(hits[:top_k]) diff --git a/extra/radeon_forge/ledger.py b/extra/radeon_forge/ledger.py new file mode 100644 index 0000000000000..fcdc8300b9c1d --- /dev/null +++ b/extra/radeon_forge/ledger.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +import json +import os +from dataclasses import asdict, is_dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Iterable + + +class ExperimentLedger: + """Append-only JSONL evidence store. + + Generated implementations are disposable; this ledger preserves hypotheses, + contracts, measurements, rejections, and approvals needed to reproduce why a + candidate was selected. + """ + + def __init__(self, path: str | os.PathLike[str]): + self.path = Path(path) + self.path.parent.mkdir(parents=True, exist_ok=True) + + @staticmethod + def _jsonable(value: Any) -> Any: + if is_dataclass(value): return asdict(value) + if isinstance(value, Path): return str(value) + if isinstance(value, dict): return {str(k): ExperimentLedger._jsonable(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): return [ExperimentLedger._jsonable(v) for v in value] + return value + + def append(self, event: str, payload: Any) -> None: + record = { + "schema_version": 1, + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "event": event, + "payload": self._jsonable(payload), + } + line = json.dumps(record, sort_keys=True, separators=(",", ":")) + "\n" + fd = os.open(self.path, os.O_APPEND | os.O_CREAT | os.O_WRONLY, 0o600) + try: + os.write(fd, line.encode("utf-8")) + os.fsync(fd) + finally: + os.close(fd) + + def records(self) -> Iterable[dict[str, Any]]: + if not self.path.exists(): return () + with self.path.open("r", encoding="utf-8") as f: + return tuple(json.loads(line) for line in f if line.strip()) diff --git a/extra/radeon_forge/oracles/__init__.py b/extra/radeon_forge/oracles/__init__.py new file mode 100644 index 0000000000000..3cad1d1b42084 --- /dev/null +++ b/extra/radeon_forge/oracles/__init__.py @@ -0,0 +1,2 @@ +from .mockgpu import MockGPUOracle, MockGPUResult +__all__ = ["MockGPUOracle", "MockGPUResult"] diff --git a/extra/radeon_forge/oracles/mockgpu.py b/extra/radeon_forge/oracles/mockgpu.py new file mode 100644 index 0000000000000..b1f9ddec36438 --- /dev/null +++ b/extra/radeon_forge/oracles/mockgpu.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +import json, os, shlex, subprocess, time +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Mapping, Sequence + + +@dataclass(frozen=True) +class MockGPUResult: + passed: bool + returncode: int + elapsed_ms: float + command: tuple[str, ...] + stdout: str + stderr: str + target: str = "gfx1100" + backend: str = "DEV=MOCK+AMD" + + def to_dict(self): return asdict(self) + + +class MockGPUOracle: + """Runs candidate correctness checks through tinygrad's integrated RDNA3 mock GPU. + + This is a semantic/instruction support gate, never a performance oracle. + Real W7900 execution remains authoritative for latency and resource behavior. + """ + def __init__(self, root: str | Path, timeout_seconds: int = 600): + self.root = Path(root).resolve() + self.timeout_seconds = timeout_seconds + + def run(self, command: Sequence[str] | str, env: Mapping[str, str] | None = None) -> MockGPUResult: + argv = tuple(shlex.split(command) if isinstance(command, str) else (str(x) for x in command)) + if not argv: raise ValueError("command must not be empty") + child_env = os.environ.copy() + child_env.update({"DEV": "MOCK+AMD", "PYTHONUNBUFFERED": "1"}) + child_env.update(env or {}) + started = time.perf_counter_ns() + proc = subprocess.run(argv, cwd=self.root, env=child_env, text=True, capture_output=True, timeout=self.timeout_seconds) + return MockGPUResult(proc.returncode == 0, proc.returncode, (time.perf_counter_ns() - started) / 1e6, argv, + proc.stdout[-50000:], proc.stderr[-50000:]) + + def run_json(self, command: Sequence[str] | str, env: Mapping[str, str] | None = None) -> str: + return json.dumps(self.run(command, env).to_dict(), default=str) diff --git a/extra/radeon_forge/permissions.py b/extra/radeon_forge/permissions.py new file mode 100644 index 0000000000000..e1e7e8b68d84d --- /dev/null +++ b/extra/radeon_forge/permissions.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import secrets +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from enum import Enum +from typing import Iterable + + +class Action(str, Enum): + INSPECT = "inspect" + COMPILE = "compile" + BENCHMARK = "benchmark" + WRITE_GENERATED_SOURCE = "write_generated_source" + RESTART_LOCAL_SERVICE = "restart_local_service" + DEPLOY = "deploy" + DOWNLOAD_MODEL = "download_model" + + +@dataclass(frozen=True) +class PermissionGrant: + token: str + actions: frozenset[Action] + reason: str + issued_at_utc: datetime + expires_at_utc: datetime + max_uses: int + + +class PermissionDenied(RuntimeError): pass + + +class PermissionController: + """Deny-by-default session permission controller. + + Grants are explicit, scoped, expiring, and usage-limited. The controller is + intentionally independent of the language model; an agent cannot authorize + its own consequential action. + """ + + def __init__(self): + self._grants: dict[str, PermissionGrant] = {} + self._uses: dict[str, int] = {} + + def issue(self, actions: Iterable[Action], reason: str, ttl_seconds: int = 900, max_uses: int = 1) -> PermissionGrant: + scope = frozenset(actions) + if not scope: raise ValueError("a permission grant must include at least one action") + if not reason.strip(): raise ValueError("a permission grant requires a reason") + if ttl_seconds <= 0 or max_uses <= 0: raise ValueError("ttl_seconds and max_uses must be positive") + now = datetime.now(timezone.utc) + grant = PermissionGrant(secrets.token_urlsafe(24), scope, reason.strip(), now, now + timedelta(seconds=ttl_seconds), max_uses) + self._grants[grant.token] = grant + self._uses[grant.token] = 0 + return grant + + def revoke(self, token: str) -> None: + self._grants.pop(token, None) + self._uses.pop(token, None) + + def authorize(self, token: str | None, action: Action) -> PermissionGrant: + if token is None or token not in self._grants: raise PermissionDenied(f"no active grant for {action.value}") + grant = self._grants[token] + if datetime.now(timezone.utc) >= grant.expires_at_utc: + self.revoke(token) + raise PermissionDenied(f"grant expired for {action.value}") + if action not in grant.actions: raise PermissionDenied(f"grant does not include {action.value}") + uses = self._uses[token] + if uses >= grant.max_uses: + self.revoke(token) + raise PermissionDenied(f"grant exhausted for {action.value}") + self._uses[token] = uses + 1 + if self._uses[token] >= grant.max_uses: + # The returned immutable grant remains valid as evidence, but the token + # cannot be reused for a second action. + self._grants.pop(token, None) + return grant diff --git a/extra/radeon_forge/planner.py b/extra/radeon_forge/planner.py new file mode 100644 index 0000000000000..249a162565552 --- /dev/null +++ b/extra/radeon_forge/planner.py @@ -0,0 +1,125 @@ +from __future__ import annotations + +import ipaddress +import json +import urllib.error +import urllib.request +from dataclasses import dataclass +from typing import Any, Mapping, Sequence +from urllib.parse import urlparse + +from .contracts import WorkloadContract +from .knowledge import RetrievalHit +from .permissions import Action + + +@dataclass(frozen=True) +class PlannerDecision: + summary: str + hypothesis: str + candidate_family: str + proposed_actions: tuple[Action, ...] + benchmark_budget: int + rationale: str + unknowns: tuple[str, ...] = () + + +class LocalEndpointRequired(ValueError): pass +class PlannerProtocolError(RuntimeError): pass + + +def validate_local_endpoint(endpoint: str) -> str: + parsed = urlparse(endpoint) + if parsed.scheme != "http" or not parsed.hostname: raise LocalEndpointRequired("planner endpoint must be a local HTTP URL") + hostname = parsed.hostname.lower() + is_loopback = hostname == "localhost" + if not is_loopback: + try: is_loopback = ipaddress.ip_address(hostname).is_loopback + except ValueError: is_loopback = False + if not is_loopback: raise LocalEndpointRequired(f"remote planner endpoint is forbidden: {hostname}") + if parsed.username or parsed.password: raise LocalEndpointRequired("credentials must not be embedded in the planner URL") + return endpoint.rstrip("/") + + +class LocalPlanner: + """OpenAI-compatible planner client restricted to loopback. + + Planner output is advisory structured data. It cannot execute tools, issue + permissions, or override deterministic oracle results. + """ + + def __init__(self, endpoint: str, model: str, timeout_seconds: int = 120): + self.endpoint = validate_local_endpoint(endpoint) + self.model = model + self.timeout_seconds = timeout_seconds + + @staticmethod + def _evidence(hits: Sequence[RetrievalHit]) -> list[dict[str, Any]]: + return [{"citation": hit.citation, "score": hit.score, "text": hit.chunk.text} for hit in hits] + + def plan(self, request: str, contract: WorkloadContract, evidence: Sequence[RetrievalHit] = (), + memory: Sequence[Mapping[str, Any]] = ()) -> PlannerDecision: + system = """You are Radeon Forge, a private AMD performance-engineering planner. +Return one JSON object only. Propose falsifiable hypotheses and bounded tool actions. +Never claim a speedup before hardware measurement. Never bypass correctness, quality, +privacy, stability, resource, or permission gates. Generated code is disposable; the +contract and oracle are authoritative.""" + user = { + "request": request, + "workload_contract": { + "name": contract.name, + "target": contract.target, + "objective": contract.objective.value, + "max_abs_error": contract.max_abs_error, + "max_rel_error": contract.max_rel_error, + "max_vgprs": contract.max_vgprs, + "max_lds_bytes": contract.max_lds_bytes, + "forbid_spills": contract.forbid_spills, + "metadata": dict(contract.metadata), + }, + "local_evidence": self._evidence(evidence), + "recent_memory": list(memory), + "required_schema": { + "summary": "string", + "hypothesis": "string", + "candidate_family": "string", + "proposed_actions": [action.value for action in Action], + "benchmark_budget": "positive integer", + "rationale": "string", + "unknowns": ["string"], + }, + } + payload = { + "model": self.model, + "temperature": 0, + "messages": [{"role": "system", "content": system}, {"role": "user", "content": json.dumps(user, sort_keys=True)}], + "response_format": {"type": "json_object"}, + } + request_obj = urllib.request.Request( + self.endpoint + "/chat/completions", + data=json.dumps(payload).encode("utf-8"), + headers={"Content-Type": "application/json"}, + method="POST", + ) + try: + with urllib.request.urlopen(request_obj, timeout=self.timeout_seconds) as response: + result = json.loads(response.read().decode("utf-8")) + except (urllib.error.URLError, TimeoutError, json.JSONDecodeError) as exc: + raise PlannerProtocolError(f"local planner request failed: {exc}") from exc + try: + content = result["choices"][0]["message"]["content"] + data = json.loads(content) + actions = tuple(Action(value) for value in data["proposed_actions"]) + budget = int(data["benchmark_budget"]) + if budget <= 0: raise ValueError("benchmark_budget must be positive") + return PlannerDecision( + summary=str(data["summary"]), + hypothesis=str(data["hypothesis"]), + candidate_family=str(data["candidate_family"]), + proposed_actions=actions, + benchmark_budget=budget, + rationale=str(data["rationale"]), + unknowns=tuple(str(value) for value in data.get("unknowns", ())), + ) + except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: + raise PlannerProtocolError(f"invalid planner response: {exc}") from exc diff --git a/extra/radeon_forge/profiling/__init__.py b/extra/radeon_forge/profiling/__init__.py new file mode 100644 index 0000000000000..d28d85b81703c --- /dev/null +++ b/extra/radeon_forge/profiling/__init__.py @@ -0,0 +1,2 @@ +from .report import Finding, build_profile_report +__all__ = ["Finding", "build_profile_report"] diff --git a/extra/radeon_forge/profiling/capture.py b/extra/radeon_forge/profiling/capture.py new file mode 100644 index 0000000000000..887f4dc4608b0 --- /dev/null +++ b/extra/radeon_forge/profiling/capture.py @@ -0,0 +1,177 @@ +from __future__ import annotations + +import csv, hashlib, json, os, shutil, subprocess, time, uuid +from dataclasses import asdict, dataclass, field +from enum import Enum +from pathlib import Path +from typing import Any, Mapping, Sequence + + +class CaptureKind(str, Enum): + COUNTERS = "counters" + ATT = "att" + + +@dataclass(frozen=True) +class CaptureRequest: + kind: CaptureKind + command: tuple[str, ...] + stage: str + session_id: str = "" + trace_id: str = "" + agent_step: int = 0 + kernel_regex: str = "" + counters: tuple[str, ...] = () + timeout_seconds: int = 900 + metadata: Mapping[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class CaptureArtifact: + path: str + size_bytes: int + sha256: str + suffix: str + + +@dataclass(frozen=True) +class CaptureResult: + capture_id: str + kind: str + passed: bool + output_directory: str + profiler_command: tuple[str, ...] + target_command: tuple[str, ...] + returncode: int | None + elapsed_ms: float + stdout_tail: str + stderr_tail: str + artifacts: tuple[CaptureArtifact, ...] + summary: Mapping[str, Any] + error: str = "" + + def to_dict(self) -> dict[str, Any]: return asdict(self) + + +class Rocprofv3Adapter: + """Evidence-preserving rocprofv3 adapter for counters and ATT/SQTT. + + CLI options are detected from the installed profiler help rather than assumed + from a particular ROCm release. Performance is always measured by the target + hardware; this adapter only captures and normalizes evidence. + """ + def __init__(self, project_root: str | Path, evidence_root: str | Path): + self.project_root = Path(project_root).resolve() + self.evidence_root = Path(evidence_root).resolve() + self.evidence_root.mkdir(parents=True, exist_ok=True) + + @staticmethod + def executable() -> str | None: return shutil.which("rocprofv3") + + def probe(self) -> dict[str, Any]: + executable = self.executable() + if executable is None: return {"available":False, "reason":"rocprofv3 is not on PATH"} + version = subprocess.run([executable, "--version"], text=True, capture_output=True, timeout=20) + help_run = subprocess.run([executable, "--help"], text=True, capture_output=True, timeout=20) + help_text = help_run.stdout + "\n" + help_run.stderr + avail_exe = shutil.which("rocprofv3-avail") + available = None + if avail_exe: + proc = subprocess.run([avail_exe, "list"], text=True, capture_output=True, timeout=60) + available = {"returncode":proc.returncode, "stdout":proc.stdout[-50000:], "stderr":proc.stderr[-10000:]} + return {"available":True, "executable":executable, "version_returncode":version.returncode, + "version":(version.stdout+version.stderr).strip(), "supports_att":"--att" in help_text, + "supports_kernel_regex":"--kernel-include-regex" in help_text, "supports_pmc":"--pmc" in help_text, + "supports_output_directory":"--output-directory" in help_text or " -d" in help_text, + "available_components":available} + + def _validate_target(self, command: Sequence[str]) -> tuple[str, ...]: + argv = tuple(str(x) for x in command) + if not argv: raise ValueError("profile target command must not be empty") + executable = Path(argv[0]).name + allowed = {"python", "python3", "pytest", "tinygrad", "radeon-forge"} + if executable not in allowed: raise ValueError(f"profile target executable {executable!r} is not allowlisted") + if any("\x00" in part for part in argv): raise ValueError("profile command contains NUL") + return argv + + def _profiler_command(self, request: CaptureRequest, output: Path, help_text: str) -> tuple[str, ...]: + executable = self.executable() + if executable is None: raise RuntimeError("rocprofv3 is not installed") + args: list[str] = [executable] + if request.kind is CaptureKind.ATT: + if "--att" not in help_text: raise RuntimeError("installed rocprofv3 does not advertise ATT support") + args.append("--att") + if "--att-simd-select" in help_text: args += ["--att-simd-select", "0x0"] + if request.kernel_regex: + if "--kernel-include-regex" not in help_text: raise RuntimeError("installed rocprofv3 cannot filter ATT by kernel regex") + args += ["--kernel-include-regex", request.kernel_regex] + elif request.kind is CaptureKind.COUNTERS: + if not request.counters: raise ValueError("counter capture requires at least one counter") + if "--pmc" not in help_text: raise RuntimeError("installed rocprofv3 does not advertise --pmc counter collection") + args += ["--pmc", ",".join(request.counters)] + else: raise ValueError(request.kind) + if "--output-directory" in help_text: args += ["--output-directory", str(output)] + else: args += ["-d", str(output)] + return tuple(args + ["--"] + list(request.command)) + + @staticmethod + def _artifacts(root: Path) -> tuple[CaptureArtifact, ...]: + ret = [] + for path in sorted(root.rglob("*")): + if not path.is_file(): continue + data = path.read_bytes() + ret.append(CaptureArtifact(str(path), len(data), hashlib.sha256(data).hexdigest(), path.suffix.lower())) + return tuple(ret) + + @staticmethod + def _csv_summary(root: Path) -> dict[str, Any]: + files, rows, columns = 0, 0, set() + numeric: dict[str, list[float]] = {} + for path in root.rglob("*.csv"): + files += 1 + try: + with path.open(newline="", encoding="utf-8", errors="replace") as handle: + reader = csv.DictReader(handle) + columns.update(reader.fieldnames or ()) + for row in reader: + rows += 1 + for key, value in row.items(): + try: numeric.setdefault(str(key), []).append(float(value)) + except (TypeError, ValueError): pass + except OSError: continue + aggregate = {key:{"count":len(values), "min":min(values), "max":max(values), "mean":sum(values)/len(values)} + for key, values in numeric.items() if values} + return {"csv_files":files, "csv_rows":rows, "columns":sorted(columns), "numeric_columns":aggregate} + + def capture(self, request: CaptureRequest) -> CaptureResult: + target = self._validate_target(request.command) + capture_id = f"{request.kind.value}-{time.strftime('%Y%m%dT%H%M%S')}-{uuid.uuid4().hex[:10]}" + output = self.evidence_root / capture_id + output.mkdir(parents=True, exist_ok=False) + manifest_path = output / "forge_capture_manifest.json" + started = time.perf_counter_ns() + returncode: int | None = None + stdout = stderr = error = "" + profiler_command: tuple[str, ...] = () + try: + executable = self.executable() + if executable is None: raise RuntimeError("rocprofv3 is not installed") + help_run = subprocess.run([executable, "--help"], text=True, capture_output=True, timeout=20) + help_text = help_run.stdout + "\n" + help_run.stderr + profiler_command = self._profiler_command(request, output, help_text) + env = {**os.environ, "PYTHONUNBUFFERED":"1", "RADEON_FORGE_CAPTURE_ID":capture_id, + "RADEON_FORGE_EXECUTION_STAGE":request.stage, "RADEON_FORGE_TRACE_ID":request.trace_id} + proc = subprocess.run(profiler_command, cwd=self.project_root, env=env, text=True, capture_output=True, + timeout=max(1, min(request.timeout_seconds, 3600))) + returncode, stdout, stderr = proc.returncode, proc.stdout, proc.stderr + except Exception as exc: error = f"{type(exc).__name__}: {exc}" + elapsed_ms = (time.perf_counter_ns() - started) / 1e6 + artifacts = tuple(x for x in self._artifacts(output) if Path(x.path) != manifest_path) + summary = {"stage":request.stage, "session_id":request.session_id, "trace_id":request.trace_id, + "agent_step":request.agent_step, "kernel_regex":request.kernel_regex, "counters":list(request.counters), + "capture_metadata":dict(request.metadata), "parsed_csv":self._csv_summary(output)} + passed = returncode == 0 and not error and bool(artifacts) + result = CaptureResult(capture_id, request.kind.value, passed, str(output), profiler_command, target, returncode, + elapsed_ms, stdout[-50000:], stderr[-50000:], artifacts, summary, error) + manifest_path.write_text(json.dumps(result.to_dict(), indent=2, default=str), encoding="utf-8") + return result diff --git a/extra/radeon_forge/profiling/compare.py b/extra/radeon_forge/profiling/compare.py new file mode 100644 index 0000000000000..8b271baa25e34 --- /dev/null +++ b/extra/radeon_forge/profiling/compare.py @@ -0,0 +1,173 @@ +from __future__ import annotations + +import csv, json, math +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Iterable, Mapping + + +@dataclass(frozen=True) +class NumericDelta: + baseline: float + candidate: float + absolute: float + relative: float | None + + +@dataclass(frozen=True) +class KernelDelta: + kernel: str + metric: str + baseline_count: int + candidate_count: int + baseline_mean: float + candidate_mean: float + absolute: float + relative: float | None + + +class IncomparableCaptures(ValueError): pass + + +def _manifest(value: str | Path | Mapping[str, Any]) -> tuple[dict[str, Any], Path | None]: + if isinstance(value, Mapping): return dict(value), None + path = Path(value) + if path.is_dir(): path = path / "forge_capture_manifest.json" + if not path.is_file(): raise FileNotFoundError(path) + payload = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(payload, dict): raise ValueError(f"capture manifest {path} is not an object") + return payload, path.parent + + +def _normalized_command(value: Any) -> tuple[str, ...]: + if not isinstance(value, list): return () + return tuple(str(x) for x in value) + + +def _counter_set(manifest: Mapping[str, Any]) -> tuple[str, ...]: + summary = manifest.get("summary", {}) + if not isinstance(summary, Mapping): return () + counters = summary.get("counters", ()) + return tuple(sorted(str(x) for x in counters)) if isinstance(counters, list) else () + + +def _workload_identity(manifest: Mapping[str, Any]) -> str: + summary = manifest.get("summary", {}) + metadata = summary.get("capture_metadata", {}) if isinstance(summary, Mapping) else {} + if isinstance(metadata, Mapping): + for key in ("workload_hash", "task_suite_hash", "input_hash", "workload_id"): + if metadata.get(key): return f"{key}:{metadata[key]}" + command = _normalized_command(manifest.get("target_command")) + return "command:" + json.dumps(command) + + +def compatibility_report(baseline: Mapping[str, Any], candidate: Mapping[str, Any]) -> dict[str, Any]: + reasons: list[str] = [] + base_summary = baseline.get("summary", {}) if isinstance(baseline.get("summary", {}), Mapping) else {} + cand_summary = candidate.get("summary", {}) if isinstance(candidate.get("summary", {}), Mapping) else {} + checks = { + "kind": (baseline.get("kind"), candidate.get("kind")), + "stage": (base_summary.get("stage"), cand_summary.get("stage")), + "target_command": (_normalized_command(baseline.get("target_command")), _normalized_command(candidate.get("target_command"))), + "workload_identity": (_workload_identity(baseline), _workload_identity(candidate)), + } + if baseline.get("kind") == "counters" or candidate.get("kind") == "counters": + checks["counters"] = (_counter_set(baseline), _counter_set(candidate)) + for name, (left, right) in checks.items(): + if left != right: reasons.append(f"{name} differs: baseline={left!r} candidate={right!r}") + if not baseline.get("passed"): reasons.append("baseline capture did not pass") + if not candidate.get("passed"): reasons.append("candidate capture did not pass") + return {"comparable": not reasons, "reasons": reasons, "checks": checks} + + +def _relative(base: float, candidate: float) -> float | None: + if base == 0: return None + return (candidate - base) / base + + +def _numeric_summary(manifest: Mapping[str, Any]) -> Mapping[str, Any]: + summary = manifest.get("summary", {}) + parsed = summary.get("parsed_csv", {}) if isinstance(summary, Mapping) else {} + numeric = parsed.get("numeric_columns", {}) if isinstance(parsed, Mapping) else {} + return numeric if isinstance(numeric, Mapping) else {} + + +def _numeric_deltas(baseline: Mapping[str, Any], candidate: Mapping[str, Any]) -> dict[str, dict[str, Any]]: + left, right = _numeric_summary(baseline), _numeric_summary(candidate) + ret: dict[str, dict[str, Any]] = {} + for key in sorted(set(left) & set(right)): + if not isinstance(left[key], Mapping) or not isinstance(right[key], Mapping): continue + if "mean" not in left[key] or "mean" not in right[key]: continue + base, cand = float(left[key]["mean"]), float(right[key]["mean"]) + ret[str(key)] = asdict(NumericDelta(base, cand, cand-base, _relative(base, cand))) + return ret + + +def _artifact_paths(manifest: Mapping[str, Any], root: Path | None, suffix: str) -> list[Path]: + paths = [] + for item in manifest.get("artifacts", []): + if not isinstance(item, Mapping) or str(item.get("suffix", "")).lower() != suffix: continue + path = Path(str(item.get("path", ""))) + if not path.is_absolute() and root is not None: path = root / path + if path.is_file(): paths.append(path) + return paths + + +def _kernel_column(fieldnames: Iterable[str]) -> str | None: + names = list(fieldnames) + preferred = ("Kernel_Name", "KernelName", "Kernel", "Name", "kernel_name", "kernel") + return next((name for name in preferred if name in names), None) + + +def _kernel_metrics(manifest: Mapping[str, Any], root: Path | None) -> dict[tuple[str, str], list[float]]: + grouped: dict[tuple[str, str], list[float]] = {} + for path in _artifact_paths(manifest, root, ".csv"): + try: + with path.open(newline="", encoding="utf-8", errors="replace") as handle: + reader = csv.DictReader(handle) + key_column = _kernel_column(reader.fieldnames or ()) + if key_column is None: continue + for row in reader: + kernel = str(row.get(key_column, "")).strip() + if not kernel: continue + for key, value in row.items(): + if key == key_column: continue + try: number = float(value) + except (TypeError, ValueError): continue + if math.isfinite(number): grouped.setdefault((kernel, str(key)), []).append(number) + except OSError: continue + return grouped + + +def _kernel_deltas(baseline: Mapping[str, Any], base_root: Path | None, + candidate: Mapping[str, Any], cand_root: Path | None) -> list[dict[str, Any]]: + left, right = _kernel_metrics(baseline, base_root), _kernel_metrics(candidate, cand_root) + ret = [] + for kernel, metric in sorted(set(left) & set(right)): + base_values, cand_values = left[(kernel, metric)], right[(kernel, metric)] + base_mean, cand_mean = sum(base_values)/len(base_values), sum(cand_values)/len(cand_values) + ret.append(asdict(KernelDelta(kernel, metric, len(base_values), len(cand_values), base_mean, cand_mean, + cand_mean-base_mean, _relative(base_mean, cand_mean)))) + return ret + + +def compare_captures(baseline_value: str | Path | Mapping[str, Any], + candidate_value: str | Path | Mapping[str, Any]) -> dict[str, Any]: + baseline, base_root = _manifest(baseline_value) + candidate, cand_root = _manifest(candidate_value) + compatibility = compatibility_report(baseline, candidate) + if not compatibility["comparable"]: raise IncomparableCaptures("; ".join(compatibility["reasons"])) + return { + "baseline_capture_id": baseline.get("capture_id"), + "candidate_capture_id": candidate.get("capture_id"), + "kind": baseline.get("kind"), + "stage": (baseline.get("summary", {}) or {}).get("stage"), + "compatibility": compatibility, + "numeric_column_deltas": _numeric_deltas(baseline, candidate), + "per_kernel_deltas": _kernel_deltas(baseline, base_root, candidate, cand_root), + "interpretation_contract": { + "lower_is_better_metrics": [], + "higher_is_better_metrics": [], + "note": "Deltas are descriptive only. Metric direction and causal meaning must come from the profiler schema or domain knowledge, not guessed from column names." + } + } diff --git a/extra/radeon_forge/profiling/report.py b/extra/radeon_forge/profiling/report.py new file mode 100644 index 0000000000000..b77d32a6e9e76 --- /dev/null +++ b/extra/radeon_forge/profiling/report.py @@ -0,0 +1,163 @@ +from __future__ import annotations + +import math, statistics +from collections import defaultdict +from dataclasses import asdict, dataclass +from typing import Any, Iterable + +from ..runtime.events import TraceEvent + + +@dataclass(frozen=True) +class Finding: + severity: str + status: str + title: str + evidence: tuple[str, ...] + recommendation: str + + +def _percentile(values: list[float], percentile: float) -> float | None: + if not values: return None + ordered = sorted(values) + idx = min(len(ordered) - 1, max(0, math.ceil(percentile * len(ordered)) - 1)) + return ordered[idx] + + +def _number(value: Any, default: float = 0.0) -> float: + try: return float(value) + except (TypeError, ValueError): return default + + +def _rows(grouped: dict[str, list[float]]) -> tuple[list[dict[str, Any]], float]: + total = sum(sum(values) for values in grouped.values()) + rows = [{"name": name, "calls": len(values), "total_ms": sum(values), "mean_ms": statistics.mean(values), + "p95_ms": _percentile(values, 0.95), "share": sum(values) / max(total, 1e-12)} for name, values in grouped.items()] + rows.sort(key=lambda row: (-row["total_ms"], row["name"])) + return rows, total + + +def _kernel_summary(events: list[TraceEvent]) -> tuple[list[dict[str, Any]], float, dict[str, int], dict[str, list[dict[str, Any]]]]: + grouped: dict[str, list[float]] = defaultdict(list) + by_stage: dict[str, dict[str, list[float]]] = defaultdict(lambda: defaultdict(list)) + stage_counts: dict[str, int] = defaultdict(int) + for event in events: + if event.kind != "kernel": continue + duration = _number(event.attributes.get("duration_ms"), event.duration_ms or 0.0) + if duration <= 0: continue + name = str(event.name or event.attributes.get("name") or event.attributes.get("kernel_name") or "unknown_kernel") + stage = str(event.attributes.get("stage", "unknown")) + grouped[name].append(duration) + by_stage[stage][name].append(duration) + stage_counts[stage] += 1 + rows, total = _rows(grouped) + stage_rows = {stage: _rows(values)[0] for stage, values in by_stage.items()} + return rows, total, dict(stage_counts), stage_rows + + +def build_profile_report(events: Iterable[TraceEvent]) -> dict[str, Any]: + evs = list(events) + durations = {kind: sum(x.duration_ms or 0.0 for x in evs if x.kind == kind) for kind in {x.kind for x in evs}} + token_events = [x for x in evs if x.name == "token"] + token_wall = [_number(x.attributes.get("wall_ms")) for x in token_events if _number(x.attributes.get("wall_ms")) > 0] + token_gpu = [_number(x.attributes.get("gpu_ms")) for x in token_events if _number(x.attributes.get("gpu_ms")) > 0] + kernels_per_token = [int(_number(x.attributes.get("kernel_count"))) for x in token_events] + token_wall_by_stage: dict[str, list[float]] = defaultdict(list) + for event in token_events: + wall = _number(event.attributes.get("wall_ms")) + if wall > 0: token_wall_by_stage[str(event.attributes.get("stage", "decode"))].append(wall) + expected_decode_profile = sum(int(_number(x.attributes.get("profile_kernel_events"))) for x in token_events) + prefill = [x for x in evs if x.name == "prefill"] + prefill_wall = sum(_number(x.attributes.get("wall_ms")) for x in prefill) + expected_prefill_profile = sum(int(_number(x.attributes.get("profile_kernel_events"))) for x in prefill) + tools = [x for x in evs if x.kind == "tool" and x.duration_ms is not None and x.name != "tool_result"] + metric_events = [x for x in evs if x.name == "metric"] + truncated = [x for x in metric_events if x.attributes.get("name") == "profile_truncated"] + hook_events = [x for x in evs if x.name == "hook"] + hook_rollbacks = [x for x in hook_events if x.attributes.get("rolled_back")] + top_kernels, profiled_kernel_ms, stage_kernel_counts, top_kernels_by_stage = _kernel_summary(evs) + captured_profile = sum(stage_kernel_counts.values()) + findings: list[Finding] = [] + + if kernels_per_token and statistics.mean(kernels_per_token) >= 20: + findings.append(Finding("high", "observed", "High kernel-launch count per decoded token", + (f"mean launches/token={statistics.mean(kernels_per_token):.1f}", f"profiled token-stage ranges={sum(v for k,v in stage_kernel_counts.items() if k != 'prefill')}"), + "Inspect the dominant launch sequence separately for first-token and steady decode, then test a fused subgraph or persistent megakernel.")) + if token_wall and token_gpu and sum(token_gpu) / max(sum(token_wall), 1e-9) < 0.65: + findings.append(Finding("high", "inferred", "GPU execution explains only part of decode wall time", + (f"GPU/wall ratio={sum(token_gpu)/sum(token_wall):.2f}",), + "Profile CPU submission, synchronization, sampling and launch bubbles before optimizing arithmetic throughput.")) + if prefill and prefill_wall > sum(token_wall): + findings.append(Finding("medium", "observed", "Prefill dominates this agent turn", + (f"prefill_ms={prefill_wall:.2f}", f"token_generation_ms={sum(token_wall):.2f}"), + "Use a prefill-specific recipe: increase stable-prefix reuse, compact repeated tool schemas, or tune prefill kernels independently of decode.")) + if tools and sum(x.duration_ms or 0 for x in tools) > sum(token_wall): + findings.append(Finding("medium", "observed", "Tool execution dominates model decode", + (f"tool_ms={sum(x.duration_ms or 0 for x in tools):.2f}",), + "Use an agent-loop or tool-execution-stage optimization; a faster model kernel alone will not materially improve task latency.")) + + first_token = token_wall_by_stage.get("first_token", []) + steady_decode = token_wall_by_stage.get("decode", []) + if first_token and steady_decode and statistics.median(first_token) > statistics.median(steady_decode) * 1.4: + findings.append(Finding("medium", "observed", "First-token and steady-decode latency are materially different", + (f"first_token_p50={statistics.median(first_token):.2f}ms", f"steady_decode_p50={statistics.median(steady_decode):.2f}ms"), + "Do not force one kernel schedule across both states. Keep separate first-token and steady-decode hook predicates and validate each independently.")) + + stage_dominants = {stage: rows[0] for stage, rows in top_kernels_by_stage.items() if rows} + distinct_dominants = {row["name"] for row in stage_dominants.values()} + if len(distinct_dominants) > 1: + findings.append(Finding("high", "observed", "Different execution stages have different dominant kernels", + tuple(f"{stage}: {row['name']} ({row['share']*100:.1f}% of captured {stage} GPU time)" for stage, row in sorted(stage_dominants.items())), + "Create separate stage-scoped optimization recipes rather than a single global replacement. Preserve a shared oracle but tune each state independently.")) + + if top_kernels: + dominant = top_kernels[0] + if dominant["share"] >= 0.25: + findings.append(Finding("high", "observed", "One kernel family dominates captured GPU time overall", + (f"kernel={dominant['name']}", f"share={dominant['share']*100:.1f}%", f"calls={dominant['calls']}", + f"total_ms={dominant['total_ms']:.2f}"), + "Map this kernel to both its model subgraph and execution stage, then generate a bounded structural alternative and retune it.")) + tiny = [row for row in top_kernels if row["mean_ms"] < 0.05] + tiny_calls = sum(row["calls"] for row in tiny) + if captured_profile and tiny_calls / max(captured_profile, 1) >= 0.35: + findings.append(Finding("medium", "inferred", "Captured execution is fragmented across many very small kernels", + (f"sub-50us calls={tiny_calls}", f"captured kernel calls={captured_profile}"), + "Inspect adjacency and materialization boundaries within the same execution stage. Fusion is promising only where the numerical oracle covers the combined subgraph.")) + else: + findings.append(Finding("info", "unknown", "No kernel-level trace is attached", + ("Only agent/model aggregate counters are available.",), + "Enable the tinygrad PROFILE event adapter or capture ROCm/SQTT evidence before making a causal hardware diagnosis.")) + + if hook_rollbacks: + findings.append(Finding("high", "observed", "An execution-stage optimization was rolled back", + tuple(f"stage={x.attributes.get('stage')} error={x.attributes.get('error')}" for x in hook_rollbacks[-4:]), + "Keep the trusted baseline active, mark the candidate failed for this execution state, and regenerate from the preserved intent and oracle.")) + + expected_profile = expected_decode_profile + expected_prefill_profile + if truncated: + findings.append(Finding("medium", "observed", "Kernel trace was truncated", + tuple(f"{x.attributes.get('stage')} captured={x.attributes.get('captured')} available={x.attributes.get('available')}" for x in truncated[:4]), + "Reduce the profiled workload or raise the capture budget before using kernel shares as complete attribution.")) + elif expected_profile and captured_profile < expected_profile: + findings.append(Finding("info", "unknown", "Some backend profile ranges were not attached to the unified trace", + (f"backend reported={expected_profile}", f"trace captured={captured_profile}"), + "Treat per-kernel attribution as partial until the transport discrepancy is resolved.")) + + stage_latency = {stage: {"count": len(values), "p50_ms": statistics.median(values) if values else None, + "p95_ms": _percentile(values, 0.95)} for stage, values in token_wall_by_stage.items()} + return { + "summary": { + "event_count": len(evs), "durations_ms_by_kind": durations, "decode_tokens": len(token_events), + "token_wall_ms_p50": statistics.median(token_wall) if token_wall else None, + "token_wall_ms_p95": _percentile(token_wall, 0.95), + "token_gpu_ms_p50": statistics.median(token_gpu) if token_gpu else None, + "token_latency_by_stage": stage_latency, + "mean_kernel_count_per_token": statistics.mean(kernels_per_token) if kernels_per_token else None, + "tool_time_ms": sum(x.duration_ms or 0 for x in tools), "prefill_wall_ms": prefill_wall, + "profiled_kernel_time_ms": profiled_kernel_ms, "profiled_kernel_calls_by_stage": stage_kernel_counts, + "kernel_profile_complete": bool(captured_profile and not truncated and captured_profile >= expected_profile), + "top_kernels": top_kernels[:20], "top_kernels_by_stage": {key: value[:10] for key, value in top_kernels_by_stage.items()}, + "hook_transitions": len(hook_events), "hook_rollbacks": len(hook_rollbacks), + }, + "findings": [asdict(x) for x in findings], + } diff --git a/extra/radeon_forge/profiling/sqtt_normalize.py b/extra/radeon_forge/profiling/sqtt_normalize.py new file mode 100644 index 0000000000000..03eaf76a57040 --- /dev/null +++ b/extra/radeon_forge/profiling/sqtt_normalize.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +import csv, json, math +from collections import Counter,defaultdict +from dataclasses import asdict,dataclass +from pathlib import Path +from typing import Any,Iterable,Mapping + + +@dataclass(frozen=True) +class ParsedArtifact: + path:str + format:str + parsed:bool + event_count:int + duration_us:float + categories:Mapping[str,int] + names:Mapping[str,int] + wave_count:int|None=None + instruction_classes:Mapping[str,int]|None=None + reason:str="" + + +class UnsupportedSQTTFormat(ValueError):pass + + +def _manifest_path(value:str|Path)->Path: + path=Path(value) + if path.is_dir():path=path/"forge_capture_manifest.json" + if not path.is_file():raise FileNotFoundError(path) + return path + + +def _artifact_paths(manifest:Mapping[str,Any],root:Path)->list[Path]: + paths=[] + for item in manifest.get("artifacts",()): + if not isinstance(item,Mapping):continue + path=Path(str(item.get("path",""))) + if not path.is_absolute():path=root/path + if path.is_file():paths.append(path) + return paths + + +def _perfetto(path:Path)->ParsedArtifact: + payload=json.loads(path.read_text(encoding="utf-8",errors="replace")) + events=payload.get("traceEvents") if isinstance(payload,Mapping) else None + if not isinstance(events,list):raise UnsupportedSQTTFormat("JSON does not contain traceEvents") + categories=Counter();names=Counter();duration=0.0;wave_ids=set();instructions=Counter() + for event in events: + if not isinstance(event,Mapping):continue + category=str(event.get("cat","uncategorized"));name=str(event.get("name","unnamed")) + categories[category]+=1;names[name]+=1 + try: + value=float(event.get("dur",0.0)) + if math.isfinite(value) and value>=0:duration+=value + except (TypeError,ValueError):pass + args=event.get("args",{}) + if isinstance(args,Mapping): + for key in ("wave","wave_id","waveId"): + if key in args:wave_ids.add(str(args[key])) + for key in ("instruction_class","instruction_type","inst_class","opcode_class"): + if key in args:instructions[str(args[key])]+=1 + return ParsedArtifact(str(path),"perfetto-json",True,len(events),duration,dict(categories),dict(names), + len(wave_ids) if wave_ids else None,dict(instructions) if instructions else None) + + +def _column(fieldnames:Iterable[str],candidates:tuple[str,...])->str|None: + names=list(fieldnames) + lowered={name.lower():name for name in names} + for candidate in candidates: + if candidate.lower() in lowered:return lowered[candidate.lower()] + return None + + +def _table(path:Path)->ParsedArtifact: + delimiter="\t" if path.suffix.lower() in {".tsv",".tab"} else "," + with path.open(newline="",encoding="utf-8",errors="replace") as handle: + reader=csv.DictReader(handle,delimiter=delimiter);fields=reader.fieldnames or [] + if not fields:raise UnsupportedSQTTFormat("table has no header") + wave_col=_column(fields,("wave","wave_id","waveid")) + inst_col=_column(fields,("instruction_class","instruction_type","inst_class","opcode_class","instruction")) + name_col=_column(fields,("kernel_name","kernel","name","event")) + category_col=_column(fields,("category","cat","type","event_type")) + duration_col=_column(fields,("duration_us","duration","dur","elapsed_us")) + start_col=_column(fields,("start_us","start","begin"));end_col=_column(fields,("end_us","end","finish")) + rows=0;duration=0.0;waves=set();instructions=Counter();names=Counter();categories=Counter() + for row in reader: + rows+=1 + if wave_col and row.get(wave_col):waves.add(str(row[wave_col])) + if inst_col and row.get(inst_col):instructions[str(row[inst_col])]+=1 + if name_col and row.get(name_col):names[str(row[name_col])]+=1 + if category_col and row.get(category_col):categories[str(row[category_col])]+=1 + try: + if duration_col and row.get(duration_col) not in (None,""):value=float(row[duration_col]) + elif start_col and end_col:value=float(row[end_col])-float(row[start_col]) + else:value=0.0 + if math.isfinite(value) and value>=0:duration+=value + except (TypeError,ValueError):pass + recognized=bool(wave_col or inst_col or duration_col or (start_col and end_col)) + if not recognized:raise UnsupportedSQTTFormat("table lacks wave, instruction or timeline columns") + return ParsedArtifact(str(path),"sqtt-table",True,rows,duration,dict(categories),dict(names), + len(waves) if waves else None,dict(instructions) if instructions else None) + + +def normalize_capture(value:str|Path)->dict[str,Any]: + manifest_path=_manifest_path(value);root=manifest_path.parent + manifest=json.loads(manifest_path.read_text(encoding="utf-8")) + parsed=[];unsupported=[] + for path in _artifact_paths(manifest,root): + if path==manifest_path:continue + try: + if path.suffix.lower()==".json":item=_perfetto(path) + elif path.suffix.lower() in {".csv",".tsv",".tab"}:item=_table(path) + else:raise UnsupportedSQTTFormat(f"no registered decoder for {path.suffix or 'binary'}") + parsed.append(item) + except (UnsupportedSQTTFormat,json.JSONDecodeError,csv.Error) as exc: + unsupported.append(ParsedArtifact(str(path),"unknown",False,0,0.0,{}, {},reason=str(exc))) + total_events=sum(item.event_count for item in parsed) + total_duration=sum(item.duration_us for item in parsed) + waves=sum(item.wave_count or 0 for item in parsed) + instructions=Counter() + for item in parsed:instructions.update(item.instruction_classes or {}) + return {"capture_id":manifest.get("capture_id"),"kind":manifest.get("kind"),"parsed":bool(parsed), + "recognized_artifacts":[asdict(item) for item in parsed],"unsupported_artifacts":[asdict(item) for item in unsupported], + "summary":{"recognized_count":len(parsed),"unsupported_count":len(unsupported),"event_count":total_events, + "duration_us":total_duration,"wave_count":waves or None,"instruction_classes":dict(instructions)}, + "contract":"Only recognized structured exports are normalized. Raw ATT/SQTT binaries remain immutable evidence until a matching decoder is installed."} diff --git a/extra/radeon_forge/profiling/tools.py b/extra/radeon_forge/profiling/tools.py new file mode 100644 index 0000000000000..e81bd30a4ba94 --- /dev/null +++ b/extra/radeon_forge/profiling/tools.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +import json +from dataclasses import asdict +from pathlib import Path +from typing import Any, Mapping + +from ..permissions import Action +from ..runtime.tool_context import current_tool_context +from ..runtime.tools import ToolRegistry, ToolSpec +from .capture import CaptureKind, CaptureRequest, Rocprofv3Adapter + + +class ProfilingTools: + def __init__(self, project_root: str | Path, evidence_root: str | Path): + self.adapter = Rocprofv3Adapter(project_root, evidence_root) + + @staticmethod + def _context() -> dict[str, Any]: + context = current_tool_context() + return asdict(context) if context is not None else {} + + @staticmethod + def _command(value: Any) -> tuple[str, ...]: + if not isinstance(value, list) or not value or not all(isinstance(x, str) for x in value): + raise ValueError("command must be a non-empty string array") + return tuple(value) + + def probe(self, args: Mapping[str, Any]) -> Any: return self.adapter.probe() + + def _capture(self, kind: CaptureKind, args: Mapping[str, Any]) -> Any: + context = self._context() + request = CaptureRequest( + kind=kind, + command=self._command(args.get("command")), + stage=str(args.get("stage", "unknown")), + session_id=str(context.get("session_id", args.get("session_id", ""))), + trace_id=str(context.get("trace_id", args.get("trace_id", ""))), + agent_step=int(context.get("agent_step", args.get("agent_step", 0))), + kernel_regex=str(args.get("kernel_regex", "")), + counters=tuple(str(x) for x in args.get("counters", ())), + timeout_seconds=int(args.get("timeout_seconds", 900)), + metadata={"tool_call_id":context.get("tool_call_id", ""), "tool_name":context.get("tool_name", ""), + **dict(args.get("metadata", {}))}, + ) + return self.adapter.capture(request).to_dict() + + def capture_counters(self, args: Mapping[str, Any]) -> Any: return self._capture(CaptureKind.COUNTERS, args) + def capture_att(self, args: Mapping[str, Any]) -> Any: return self._capture(CaptureKind.ATT, args) + + def list_captures(self, args: Mapping[str, Any]) -> Any: + captures = [] + for path in sorted(self.adapter.evidence_root.glob("*/forge_capture_manifest.json"), reverse=True): + try: + payload = json.loads(path.read_text(encoding="utf-8")) + captures.append({"capture_id":payload.get("capture_id", path.parent.name), "kind":payload.get("kind"), + "passed":payload.get("passed"), "elapsed_ms":payload.get("elapsed_ms"), "output_directory":str(path.parent), + "summary":payload.get("summary", {}), "artifact_count":len(payload.get("artifacts", [])), + "error":payload.get("error", "")}) + except Exception: continue + return {"captures":captures[:max(1, min(int(args.get("limit", 50)), 500))]} + + def inspect_capture(self, args: Mapping[str, Any]) -> Any: + capture_id = str(args["capture_id"]) + if "/" in capture_id or ".." in capture_id: raise ValueError("invalid capture id") + path = self.adapter.evidence_root / capture_id / "forge_capture_manifest.json" + if not path.is_file(): raise KeyError(f"unknown capture {capture_id}") + return json.loads(path.read_text(encoding="utf-8")) + + def install(self, registry: ToolRegistry) -> None: + registry.register(ToolSpec("probe_rocm_profiler", "Inspect the installed local rocprofv3/ATT capabilities without running a workload", + {"type":"object","properties":{}}), self.probe) + registry.register(ToolSpec("capture_rocm_counters", "Capture selected hardware performance counters for one local command and execution state", + {"type":"object","required":["command","stage","counters"],"properties":{ + "command":{"type":"array","items":{"type":"string"}}, "stage":{"type":"string"}, + "counters":{"type":"array","items":{"type":"string"}}, "kernel_regex":{"type":"string"}, + "timeout_seconds":{"type":"integer"}, "metadata":{"type":"object"}}}, Action.BENCHMARK), self.capture_counters) + registry.register(ToolSpec("capture_rocm_att", "Capture AMD Advanced Thread Trace/SQTT evidence for one local command and execution state", + {"type":"object","required":["command","stage"],"properties":{ + "command":{"type":"array","items":{"type":"string"}}, "stage":{"type":"string"}, + "kernel_regex":{"type":"string"}, "timeout_seconds":{"type":"integer"}, + "metadata":{"type":"object"}}}, Action.BENCHMARK), self.capture_att) + registry.register(ToolSpec("list_profile_captures", "List immutable local ROCm counter and ATT capture manifests", + {"type":"object","properties":{"limit":{"type":"integer"}}}), self.list_captures) + registry.register(ToolSpec("inspect_profile_capture", "Inspect one profile manifest, artifact hashes and normalized evidence summary", + {"type":"object","required":["capture_id"],"properties":{"capture_id":{"type":"string"}}}), self.inspect_capture) diff --git a/extra/radeon_forge/recipe_cli.py b/extra/radeon_forge/recipe_cli.py new file mode 100644 index 0000000000000..dacde5645759c --- /dev/null +++ b/extra/radeon_forge/recipe_cli.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +import argparse, json +from dataclasses import asdict +from pathlib import Path + +from .synthesis import CandidateWorkspace, ForgeRecipe, RecipeLibrary, export_recipe + + +def main() -> None: + parser = argparse.ArgumentParser(description="Import, inspect and export portable Radeon Forge optimization recipes") + sub = parser.add_subparsers(dest="command", required=True) + + inspect_p = sub.add_parser("inspect", help="Validate and inspect one .forge.toml file without installing it") + inspect_p.add_argument("recipe", type=Path) + + install_p = sub.add_parser("install", help="Install a recipe into a local Forge optimization library") + install_p.add_argument("recipe", type=Path) + install_p.add_argument("--workspace", type=Path, default=Path.home() / ".cache" / "radeon-forge") + + list_p = sub.add_parser("list", help="List installed portable recipes") + list_p.add_argument("--workspace", type=Path, default=Path.home() / ".cache" / "radeon-forge") + + export_p = sub.add_parser("export", help="Export a local contract and optional implementation cache as one file") + export_p.add_argument("--workspace", type=Path, default=Path.home() / ".cache" / "radeon-forge") + export_p.add_argument("--spec-id", required=True) + export_p.add_argument("--candidate-id") + export_p.add_argument("--output", type=Path, required=True) + + args = parser.parse_args() + if args.command == "inspect": + recipe = ForgeRecipe.load(args.recipe) + payload = asdict(recipe) + payload.pop("source_text", None) + payload["recipe_id"] = recipe.recipe_id + payload["artifacts"] = [{**asdict(x), "content": f"<{len(x.content.encode())} bytes>", "sha256": x.sha256} for x in recipe.artifacts] + print(json.dumps(payload, indent=2, default=str)) + return + + workspace = CandidateWorkspace(args.workspace) + library = RecipeLibrary(workspace) + if args.command == "install": print(json.dumps(asdict(library.install_file(args.recipe)), indent=2)) + elif args.command == "list": print(json.dumps([asdict(x) for x in library.installed()], indent=2)) + elif args.command == "export": + path = export_recipe(workspace, args.spec_id, args.output, args.candidate_id) + print(json.dumps({"path": str(path.resolve()), "bytes": path.stat().st_size}, indent=2)) + + +if __name__ == "__main__": main() diff --git a/extra/radeon_forge/recipes/batch1_decode_rdna3.forge.toml b/extra/radeon_forge/recipes/batch1_decode_rdna3.forge.toml new file mode 100644 index 0000000000000..098243619d877 --- /dev/null +++ b/extra/radeon_forge/recipes/batch1_decode_rdna3.forge.toml @@ -0,0 +1,66 @@ +format_version = 1 +name = "batch1-decode-rdna3" +description = "Portable optimization contract for low-latency, batch-one transformer decode on gfx1100." +operation = "Specialize one transformer decode block or high-impact decode subgraph for a fixed local private-agent workload." +target = "gfx1100" +objective = "Minimize P95 inter-token latency and launches per token without violating numerical, KV-cache, memory, or task-quality constraints." +agent_brief = """ +Treat this file as intent and evidence, not as a kernel DSL. Inspect the local model, shapes, generated tinygrad schedule, ISA and profile before choosing an implementation. You may fuse, reorder, specialize, replace abstractions, or generate a direct persistent implementation. The optional implementation cache is only a starting point. Keep observations separate from inferences, predict measurable changes before each experiment, and let the independent oracles reject bad ideas. +""" +invariants = [ + "Match the trusted tinygrad reference within the configured tolerance.", + "Preserve KV-cache indexing and visibility semantics across repeated decode steps.", + "Pass tinygrad DEV=MOCK+AMD before any real-hardware benchmark unless explicitly waived.", + "Never use MockGPU time as performance evidence.", + "Preserve the frozen agent workload task-success threshold.", + "Do not make outbound network calls during core inference, profiling, generation, or validation." +] +unknowns = [ + "The best fusion boundary is workload- and model-specific.", + "The optimal register/LDS tradeoff must be measured on the recipient W7900.", + "A full megakernel may lose to a smaller fused subgraph; benchmark both." +] +language = "python" +extension = ".py" +entrypoint = "build_kernel" + +[shapes] +batch = 1 +sequence = "single decode token" +model = "bound by the recipient workload" + +[dtypes] +activations = "bf16 or fp16" +accumulation = "fp32 where the oracle requires it" +kv_cache = "bound by the recipient model" + +[compatibility] +arch = "gfx1100" +device_class = "Radeon PRO W7900" +same_hardware_recommended = true + +[acceptance] +maximum_task_success_drop_percent = 1.0 +maximum_invalid_tool_calls = 0 +maximum_crashes = 0 +external_network_calls = 0 +primary_metric = "p95_inter_token_latency_ms" + +[oracle] +mockgpu_command = ["python3", "{candidate}", "--forge-mode", "mockgpu"] +hardware_command = ["python3", "{candidate}", "--forge-mode", "benchmark", "--json"] +heldout_command = ["python3", "{candidate}", "--forge-mode", "heldout", "--json"] + +[[artifact]] +path = "knowledge/optimization_notes.md" +role = "knowledge" +description = "Free-form domain knowledge for the intelligent executor; not executable code." +content = """ +# Batch-one RDNA3 decode notes + +Prioritize end-to-end decode latency over isolated FLOP/s. Useful evidence includes launch count per token, per-layer and per-kernel time, VGPR/SGPR/LDS/scratch metadata, occupancy, memory traffic, wait/stall evidence, prefix/KV-cache behavior, and tool-resumption overhead. + +Candidate ideas are not requirements: direct model-shape specialization, residual plus RMSNorm fusion, dequantization plus GEMV, RoPE/layout fusion, attention-decode specialization, persistent transformer-block execution, and removal of small launch boundaries. An intelligent agent may discard all of these when the profile contradicts them. + +The implementation is disposable. Durable knowledge is the intent, invariants, reference behavior, numerical tolerance, frozen workload, hardware evidence, failed experiments, and acceptance criteria. +""" diff --git a/extra/radeon_forge/runtime/__init__.py b/extra/radeon_forge/runtime/__init__.py new file mode 100644 index 0000000000000..682b5b0e81570 --- /dev/null +++ b/extra/radeon_forge/runtime/__init__.py @@ -0,0 +1,10 @@ +from .backend import BackendCapabilities, GenerationEvent, GenerationRequest, InferenceBackend, JsonlProcessBackend +from .engine import ForgeEngine +from .events import TraceEvent, TraceRecorder +from .jobs import JobSnapshot, LocalJobManager +from .session import AgentSession, SessionState +from .tools import ToolCall, ToolRegistry, ToolResult, ToolSpec, WorkspaceTools + +__all__ = ["AgentSession", "BackendCapabilities", "ForgeEngine", "GenerationEvent", "GenerationRequest", "InferenceBackend", + "JobSnapshot", "JsonlProcessBackend", "LocalJobManager", "SessionState", "ToolCall", "ToolRegistry", "ToolResult", + "ToolSpec", "TraceEvent", "TraceRecorder", "WorkspaceTools"] diff --git a/extra/radeon_forge/runtime/backend.py b/extra/radeon_forge/runtime/backend.py new file mode 100644 index 0000000000000..c515c5794b4f1 --- /dev/null +++ b/extra/radeon_forge/runtime/backend.py @@ -0,0 +1,115 @@ +from __future__ import annotations + +import json, queue, subprocess, threading +from dataclasses import dataclass, field +from typing import Any, Iterator, Mapping, Protocol, Sequence + + +@dataclass(frozen=True) +class BackendCapabilities: + streaming: bool = True + structured_tools: bool = False + prefix_cache: bool = False + persistent_kv: bool = False + kernel_metrics: bool = False + cancellation: bool = False + batched_prefill: bool = False + stage_hooks: bool = False + local_only: bool = True + + +@dataclass(frozen=True) +class GenerationRequest: + session_id: str + messages: Sequence[Mapping[str, Any]] + tools: Sequence[Mapping[str, Any]] = () + max_tokens: int = 256 + temperature: float = 0.0 + stop: Sequence[str] = () + metadata: Mapping[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class GenerationEvent: + kind: str + text: str = "" + tool_call: Mapping[str, Any] | None = None + finish_reason: str | None = None + metrics: Mapping[str, Any] = field(default_factory=dict) + + +class InferenceBackend(Protocol): + @property + def name(self) -> str: ... + @property + def capabilities(self) -> BackendCapabilities: ... + @property + def runtime_metadata(self) -> Mapping[str, Any]: ... + def stream(self, request: GenerationRequest) -> Iterator[GenerationEvent]: ... + def close(self) -> None: ... + + +class JsonlProcessBackend: + """Long-lived local model subprocess using a strict JSON-lines protocol.""" + def __init__(self, command: Sequence[str], env: Mapping[str, str] | None = None): + import os + child_env = os.environ.copy() + child_env.update(env or {}) + self._proc = subprocess.Popen(list(command), stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, + text=True, bufsize=1, env=child_env) + if self._proc.stdin is None or self._proc.stdout is None or self._proc.stderr is None: raise RuntimeError("failed to open model process pipes") + self._stdin, self._stdout, self._stderr = self._proc.stdin, self._proc.stdout, self._proc.stderr + self._lock = threading.RLock() + self._stderr_lines: queue.Queue[str] = queue.Queue() + threading.Thread(target=self._drain_stderr, daemon=True).start() + ready = self._stdout.readline() + if not ready: raise RuntimeError(f"model process exited during startup: {self.stderr_tail()}") + payload = json.loads(ready) + if payload.get("kind") != "ready": raise RuntimeError(f"invalid backend handshake: {payload}") + self._name = str(payload.get("name", "jsonl-local-model")) + self._capabilities = BackendCapabilities(**payload.get("capabilities", {})) + self._runtime_metadata = {str(key): value for key, value in payload.items() if key not in {"kind", "name", "capabilities"}} + + @property + def name(self) -> str: return self._name + @property + def capabilities(self) -> BackendCapabilities: return self._capabilities + @property + def runtime_metadata(self) -> Mapping[str, Any]: return dict(self._runtime_metadata) + + def _drain_stderr(self) -> None: + for line in self._stderr: + self._stderr_lines.put(line.rstrip()) + while self._stderr_lines.qsize() > 200: + try: self._stderr_lines.get_nowait() + except queue.Empty: break + + def stderr_tail(self, n: int = 20) -> str: + lines = list(self._stderr_lines.queue) + return "\n".join(lines[-n:]) + + def stream(self, request: GenerationRequest) -> Iterator[GenerationEvent]: + with self._lock: + self._stdin.write(json.dumps({"op": "generate", "request": { + "session_id": request.session_id, "messages": list(request.messages), "tools": list(request.tools), + "max_tokens": request.max_tokens, "temperature": request.temperature, "stop": list(request.stop), + "metadata": dict(request.metadata)}}) + "\n") + self._stdin.flush() + while True: + line = self._stdout.readline() + if not line: raise RuntimeError(f"model process terminated: {self.stderr_tail()}") + payload = json.loads(line) + if payload.get("kind") == "error": raise RuntimeError(str(payload.get("error", "backend error"))) + ev = GenerationEvent(kind=str(payload["kind"]), text=str(payload.get("text", "")), + tool_call=payload.get("tool_call"), finish_reason=payload.get("finish_reason"), + metrics=payload.get("metrics", {})) + yield ev + if ev.kind == "done": break + + def close(self) -> None: + if self._proc.poll() is not None: return + try: + self._stdin.write(json.dumps({"op": "shutdown"}) + "\n") + self._stdin.flush() + self._proc.wait(timeout=5) + except Exception: self._proc.kill() diff --git a/extra/radeon_forge/runtime/engine.py b/extra/radeon_forge/runtime/engine.py new file mode 100644 index 0000000000000..0dde5a45ec2f1 --- /dev/null +++ b/extra/radeon_forge/runtime/engine.py @@ -0,0 +1,155 @@ +from __future__ import annotations + +import contextlib, threading, uuid +from dataclasses import asdict +from pathlib import Path +from typing import Any, Callable, Iterator, Mapping, Sequence, TypeVar + +from ..permissions import PermissionController +from ..profiling.report import build_profile_report +from ..synthesis import CandidateWorkspace, OptimizationTools, RuntimeFingerprint, SafeHookRegistry, install_default_specs +from .backend import GenerationEvent, GenerationRequest, InferenceBackend +from .jobs import LocalJobManager +from .session import AgentSession, SessionState +from .tool_prompt import inject_tool_instruction +from .tools import ToolCall, ToolRegistry, WorkspaceTools + + +T = TypeVar("T") + +DEFAULT_SYSTEM_PROMPT = """You are Radeon Forge, a private local software and inference performance engineer. +Use evidence before making performance claims. Separate observations, inferences and unknowns. Prefer reversible changes and preserve correctness. You may inspect the private workspace, author disposable target-specific kernel candidates, import portable optimization recipes, validate implementations through tinygrad MockGPU, and request permission for real W7900 benchmarks. Never treat MockGPU timing as performance evidence. + +Optimization recipes are contracts, not programming languages. A hook has two independent coordinates: where it intercepts the stack (scheduler, model block, subgraph, kernel, KV cache, sampler) and when it applies (prefill, first token, steady decode, tool resume, KV append, sampling, or a workload predicate). Read the free-form intent, invariants, oracle, stage predicate, knowledge and failed experiments; then use your judgment to regenerate or radically restructure implementations. Cached source is merely one prior compilation and must never outrank the oracle.""" + + +class ForgeEngine: + def __init__(self, backend: InferenceBackend, workspace: str | Path, system_prompt: str = DEFAULT_SYSTEM_PROMPT): + self.backend, self.workspace, self.system_prompt = backend, Path(workspace).resolve(), system_prompt + self.permissions = PermissionController() + self.tools = ToolRegistry(self.permissions) + WorkspaceTools(self.workspace).install(self.tools) + self.optimization_workspace = CandidateWorkspace(self.workspace / ".radeon_forge") + install_default_specs(self.optimization_workspace) + self.hooks = SafeHookRegistry(self.optimization_workspace) + self.optimization_tools = OptimizationTools(self.optimization_workspace, self.workspace, self.hooks, self.runtime_fingerprint) + self.optimization_tools.install(self.tools) + self.jobs = LocalJobManager() + self._sessions: dict[str, AgentSession] = {} + self._lock = threading.RLock() + self._generation_gate = threading.Lock() + self._generation_state_lock = threading.RLock() + self._active_generation_id: str | None = None + + def runtime_fingerprint(self) -> RuntimeFingerprint: + metadata = getattr(self.backend, "runtime_metadata", {}) + if callable(metadata): metadata = metadata() + return RuntimeFingerprint.from_mapping(metadata if isinstance(metadata, Mapping) else {}) + + def inference_metadata(self) -> dict[str, Any]: + metadata = self.hooks.runtime_metadata() + for item in metadata.get("active_hooks", []): + try: + record = self.optimization_workspace.load_candidate(str(item["candidate_id"])) + item["parameters"] = dict(record.evidence.get("selected_parameters", {})) + except Exception: item["parameters"] = {} + return {**metadata, "runtime_fingerprint": asdict(self.runtime_fingerprint())} + + def _session_metadata(self, session: AgentSession, step: int) -> Mapping[str, Any]: return self.inference_metadata() + + @contextlib.contextmanager + def _generation_slot(self, generation_id: str): + with self._generation_gate: + with self._generation_state_lock: self._active_generation_id = generation_id + try: yield + finally: + with self._generation_state_lock: + if self._active_generation_id == generation_id: self._active_generation_id = None + + def _run_generation(self, generation_id: str, fn: Callable[[], T]) -> T: + with self._generation_slot(generation_id): return fn() + + @property + def active_generation_id(self) -> str | None: + with self._generation_state_lock: return self._active_generation_id + + def stream_inference(self, messages: Sequence[Mapping[str, Any]], tools: Sequence[Mapping[str, Any]] = (), + max_tokens: int = 256, temperature: float = 0.0, session_id: str | None = None, + stop: Sequence[str] = ()) -> Iterator[GenerationEvent]: + """Serve the resident local model without losing Forge hooks, tools or profiling events.""" + generation_id = session_id or uuid.uuid4().hex + rendered_messages = inject_tool_instruction(messages, tools) + request = GenerationRequest(generation_id, tuple(rendered_messages), tuple(tools), max_tokens, temperature, + tuple(stop), metadata=self.inference_metadata()) + def iterator(): + with self._generation_slot(generation_id): yield from self.backend.stream(request) + return iterator() + + def create_session(self) -> AgentSession: + with self._lock: + session = AgentSession(self.backend, self.tools, self.system_prompt, metadata_provider=self._session_metadata) + self._sessions[session.session_id] = session + return session + + def session(self, session_id: str) -> AgentSession: + try: return self._sessions[session_id] + except KeyError as exc: raise KeyError(f"unknown session {session_id}") from exc + + def sessions(self) -> list[dict[str, Any]]: + with self._lock: return [{"session_id": x.session_id, "state": x.state.value, "trace_id": x.trace.trace_id} for x in self._sessions.values()] + + def run_message(self, session_id: str, content: str, max_tokens: int = 512, temperature: float = 0.0): + session = self.session(session_id) + return self._run_generation(session_id, lambda: session.send(content, max_tokens, temperature)) + + def submit_message(self, session_id: str, content: str, max_tokens: int = 512, temperature: float = 0.0) -> dict[str, Any]: + job = self.jobs.submit("agent_turn", session_id, lambda: self.run_message(session_id, content, max_tokens, temperature)) + return asdict(job) + + def run_tool_approval(self, session_id: str, reason: str, max_tokens: int = 512): + session = self.session(session_id) + token = self.grant_for_pending_tool(session_id, reason) + return self._run_generation(session_id, lambda: session.approve_tool(token, max_tokens)) + + def submit_tool_approval(self, session_id: str, reason: str, max_tokens: int = 512) -> dict[str, Any]: + job = self.jobs.submit("tool_and_resume", session_id, lambda: self.run_tool_approval(session_id, reason, max_tokens)) + return asdict(job) + + def cancel_generation(self, session_id: str) -> dict[str, Any]: + session = self.session(session_id) + if not self.backend.capabilities.cancellation: raise RuntimeError("resident model backend does not support cooperative cancellation") + active = self.active_generation_id + if active != session_id or session.state is not SessionState.GENERATING: + raise RuntimeError(f"session {session_id} is not the active model generation") + self.backend.cancel() + return {"requested": True, "session_id": session_id, "active_generation_id": active, + "state": session.state.value, "model_remains_resident": True, "kv_state_preserved": True} + + def grant_for_pending_tool(self, session_id: str, reason: str, max_uses: int = 1) -> str: + session = self.session(session_id) + if session.pending_tool_call is None: raise RuntimeError("session has no pending tool") + action = self.tools.spec(session.pending_tool_call.name).action + return self.permissions.issue([action], reason, max_uses=max_uses).token + + def execute_explicit_ui_tool(self, name: str, arguments: Mapping[str, Any], reason: str) -> Any: + """Execute one direct UI mutation through the same scoped permission boundary as the agent.""" + spec = self.tools.spec(name) + grant = self.permissions.issue([spec.action], reason.strip() or f"Explicit local UI action: {name}", max_uses=1) + result = self.tools.execute(ToolCall(uuid.uuid4().hex, name, dict(arguments)), grant.token) + if not result.ok: + if isinstance(result.output, Mapping) and result.output.get("error"): raise RuntimeError(str(result.output["error"])) + raise RuntimeError(f"{name} failed") + return result.output + + def profile(self, session_id: str) -> dict[str, Any]: return build_profile_report(self.session(session_id).trace.events()) + + def optimization_state(self) -> dict[str, Any]: + return {"specs": [asdict(x) | {"spec_id": x.spec_id} for x in self.optimization_workspace.specs()], + "candidates": [asdict(x) for x in self.optimization_workspace.candidates()], + "recipes": [asdict(x) for x in self.optimization_tools.recipes.installed()], + "active_hooks": [asdict(x) for x in self.hooks.active()], + "runtime_fingerprint": asdict(self.runtime_fingerprint()), + "runtime_adapters": sorted(self.hooks.SUPPORTED_ADAPTERS), + "active_generation_id": self.active_generation_id} + + def close(self) -> None: self.backend.close() diff --git a/extra/radeon_forge/runtime/events.py b/extra/radeon_forge/runtime/events.py new file mode 100644 index 0000000000000..3f75c03cdfca2 --- /dev/null +++ b/extra/radeon_forge/runtime/events.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +import contextlib, json, threading, time, uuid +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any, Iterator, Mapping + + +def now_ns() -> int: return time.perf_counter_ns() + + +@dataclass(frozen=True) +class TraceEvent: + event_id: str + trace_id: str + kind: str + name: str + start_ns: int + end_ns: int | None = None + parent_id: str | None = None + attributes: Mapping[str, Any] = field(default_factory=dict) + + @property + def duration_ms(self) -> float | None: + return None if self.end_ns is None else (self.end_ns - self.start_ns) / 1e6 + + def to_dict(self) -> dict[str, Any]: + ret = asdict(self) + ret["duration_ms"] = self.duration_ms + return ret + + +class TraceRecorder: + """Thread-safe hierarchical trace recorder for agent, model, tool and GPU events.""" + def __init__(self, trace_id: str | None = None): + self.trace_id = trace_id or uuid.uuid4().hex + self._events: list[TraceEvent] = [] + self._lock = threading.RLock() + + def point(self, kind: str, event_name: str, parent_id: str | None = None, **attributes: Any) -> TraceEvent: + """Record an instantaneous event. + + `event_name` is deliberately not called `name`: backend evidence commonly + carries a `name` attribute for the concrete GPU kernel. Keeping those two + namespaces separate lets an event be named `kernel` while retaining the + real kernel symbol in its attributes. + """ + ev = TraceEvent(uuid.uuid4().hex, self.trace_id, kind, event_name, now_ns(), now_ns(), parent_id, attributes) + with self._lock: self._events.append(ev) + return ev + + def duration(self, kind: str, event_name: str, duration_ms: float, parent_id: str | None = None, **attributes: Any) -> TraceEvent: + """Record a duration measured by another clock domain. + + GPU profile timestamps are device-local and cannot be placed exactly on the + CPU wall-clock axis without calibration. We preserve their measured duration + and causal parent while anchoring the range at ingestion time. The attributes + retain source/stage/order for later calibrated timeline adapters. + """ + duration = max(0.0, float(duration_ms)) + end = now_ns() + start = end - int(duration * 1e6) + ev = TraceEvent(uuid.uuid4().hex, self.trace_id, kind, event_name, start, end, parent_id, attributes) + with self._lock: self._events.append(ev) + return ev + + @contextlib.contextmanager + def span(self, kind: str, event_name: str, parent_id: str | None = None, **attributes: Any) -> Iterator[TraceEvent]: + start = now_ns() + provisional = TraceEvent(uuid.uuid4().hex, self.trace_id, kind, event_name, start, None, parent_id, attributes) + try: yield provisional + except BaseException as exc: + attrs = dict(attributes) + attrs.update({"status": "error", "error_type": type(exc).__name__, "error": str(exc)}) + with self._lock: self._events.append(TraceEvent(provisional.event_id, self.trace_id, kind, event_name, start, now_ns(), parent_id, attrs)) + raise + else: + attrs = dict(attributes) + attrs.setdefault("status", "ok") + with self._lock: self._events.append(TraceEvent(provisional.event_id, self.trace_id, kind, event_name, start, now_ns(), parent_id, attrs)) + + def extend(self, events: list[TraceEvent]) -> None: + with self._lock: self._events.extend(events) + + def events(self) -> tuple[TraceEvent, ...]: + with self._lock: return tuple(sorted(self._events, key=lambda x: (x.start_ns, x.event_id))) + + def to_dict(self) -> dict[str, Any]: return {"trace_id": self.trace_id, "events": [x.to_dict() for x in self.events()]} + + def write_json(self, path: str | Path) -> Path: + out = Path(path) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(json.dumps(self.to_dict(), indent=2, default=str), encoding="utf-8") + return out + + def chrome_trace(self) -> dict[str, Any]: + events = [] + for ev in self.events(): + if ev.end_ns is None: continue + events.append({"name": ev.name, "cat": ev.kind, "ph": "X", "ts": ev.start_ns / 1000, + "dur": (ev.end_ns - ev.start_ns) / 1000, "pid": 1, "tid": ev.kind, + "args": dict(ev.attributes)}) + return {"traceEvents": events, "displayTimeUnit": "ms"} \ No newline at end of file diff --git a/extra/radeon_forge/runtime/jobs.py b/extra/radeon_forge/runtime/jobs.py new file mode 100644 index 0000000000000..6a76f9409a32b --- /dev/null +++ b/extra/radeon_forge/runtime/jobs.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +import threading, time, uuid +from dataclasses import asdict, dataclass +from typing import Any, Callable + + +@dataclass(frozen=True) +class JobSnapshot: + job_id: str + kind: str + state: str + session_id: str + created_at_s: float + started_at_s: float | None = None + completed_at_s: float | None = None + error: str | None = None + error_type: str | None = None + + +class LocalJobManager: + """Small in-process job manager for non-blocking local UI operations.""" + def __init__(self): + self._jobs: dict[str, JobSnapshot] = {} + self._lock = threading.RLock() + + def _set(self, snapshot: JobSnapshot) -> None: + with self._lock: self._jobs[snapshot.job_id] = snapshot + + def submit(self, kind: str, session_id: str, fn: Callable[[], Any]) -> JobSnapshot: + job_id, created = uuid.uuid4().hex, time.time() + initial = JobSnapshot(job_id, kind, "queued", session_id, created) + self._set(initial) + + def run(): + started = time.time() + self._set(JobSnapshot(job_id, kind, "running", session_id, created, started_at_s=started)) + try: + fn() + except BaseException as exc: + self._set(JobSnapshot(job_id, kind, "failed", session_id, created, started, time.time(), str(exc), type(exc).__name__)) + else: + self._set(JobSnapshot(job_id, kind, "completed", session_id, created, started, time.time())) + + threading.Thread(target=run, name=f"radeon-forge-{kind}-{job_id[:8]}", daemon=True).start() + return initial + + def snapshot(self, job_id: str) -> JobSnapshot: + with self._lock: + try: return self._jobs[job_id] + except KeyError as exc: raise KeyError(f"unknown job {job_id}") from exc + + def snapshots(self, session_id: str | None = None) -> list[dict[str, Any]]: + with self._lock: + jobs = list(self._jobs.values()) + if session_id is not None: jobs = [job for job in jobs if job.session_id == session_id] + return [asdict(job) for job in sorted(jobs, key=lambda x: x.created_at_s, reverse=True)] diff --git a/extra/radeon_forge/runtime/session.py b/extra/radeon_forge/runtime/session.py new file mode 100644 index 0000000000000..7de3e59efc2e1 --- /dev/null +++ b/extra/radeon_forge/runtime/session.py @@ -0,0 +1,229 @@ +from __future__ import annotations + +import json, threading, time, uuid +from dataclasses import asdict, dataclass, field +from enum import Enum +from typing import Any, Callable, Mapping + +from .backend import GenerationRequest, InferenceBackend +from .events import TraceRecorder +from .tools import ToolCall, ToolRegistry, parse_tool_call + + +class SessionState(str, Enum): + IDLE = "idle" + GENERATING = "generating" + AWAITING_TOOL_APPROVAL = "awaiting_tool_approval" + RUNNING_TOOL = "running_tool" + COMPLETED = "completed" + CANCELLED = "cancelled" + FAILED = "failed" + + +@dataclass(frozen=True) +class SessionEvent: + sequence: int + kind: str + data: Mapping[str, Any] + timestamp_s: float = field(default_factory=time.time) + + +class AgentSession: + """Persistent private-agent session with explicit tool approval and full tracing.""" + def __init__(self, backend: InferenceBackend, tools: ToolRegistry, system_prompt: str, session_id: str | None = None, + max_agent_steps: int = 8, trace: TraceRecorder | None = None, + metadata_provider: Callable[[AgentSession, int], Mapping[str, Any]] | None = None): + self.session_id = session_id or uuid.uuid4().hex + self.backend, self.tools, self.system_prompt = backend, tools, system_prompt + self.max_agent_steps, self.metadata_provider = max_agent_steps, metadata_provider + self.trace = trace or TraceRecorder() + self.messages: list[dict[str, Any]] = [{"role": "system", "content": self._system_prompt()}] + self.events: list[SessionEvent] = [] + self.pending_tool_call: ToolCall | None = None + self.pending_assistant_content = "" + self.partial_output = "" + self.last_finish_reason: str | None = None + self.state = SessionState.IDLE + self._sequence = 0 + self._lock = threading.RLock() + + def _system_prompt(self) -> str: + schemas = self.tools.schemas() + tool_text = "\n".join(f"- {x['function']['name']}: {x['function']['description']} schema={x['function']['parameters']}" for x in schemas) + return self.system_prompt.strip() + ("\n\nAvailable local tools:\n" + tool_text if schemas else "") + """ + +Tool protocol: when a tool is required, output exactly one object wrapped as +{"name":"tool_name","arguments":{...}} +and no surrounding prose. Tool execution always requires user approval. Never +claim a tool result before it is returned. All inference and tools are local. +""" + + @staticmethod + def _tool_content(call: ToolCall) -> str: + payload = {"id": call.call_id, "name": call.name, "arguments": dict(call.arguments)} + return f"{json.dumps(payload, separators=(',', ':'))}" + + @staticmethod + def _looks_like_partial_tool_protocol(output: str) -> bool: + stripped = output.lstrip() + marker = "" + return bool(stripped) and (marker.startswith(stripped) or stripped.startswith(marker)) + + def _emit(self, kind: str, **data: Any) -> SessionEvent: + self._sequence += 1 + event = SessionEvent(self._sequence, kind, data) + self.events.append(event) + return event + + def send(self, content: str, max_tokens: int = 512, temperature: float = 0.0) -> tuple[SessionEvent, ...]: + with self._lock: + if self.state in {SessionState.GENERATING, SessionState.RUNNING_TOOL}: raise RuntimeError("session is busy") + if self.pending_tool_call is not None: raise RuntimeError("approve or reject the pending tool call first") + text = content.strip() + if not text: raise ValueError("message must not be empty") + self.messages.append({"role": "user", "content": text}) + self._emit("message", role="user", content=text) + return self._generate(max_tokens, temperature) + + def _record_backend_event(self, event, parent_id: str) -> None: + metrics = dict(event.metrics) + self._emit(event.kind, **metrics) + if event.kind == "kernel": + name = str(metrics.pop("name", metrics.pop("kernel_name", "kernel"))) + duration_ms = float(metrics.pop("duration_ms", metrics.pop("duration_us", 0.0) / 1000.0)) + self.trace.duration("kernel", name, duration_ms, parent_id, **metrics) + elif event.kind in {"prefill", "decode"} and float(metrics.get("wall_ms", 0.0)) > 0: + self.trace.duration("inference", event.kind, float(metrics["wall_ms"]), parent_id, **metrics) + else: + if "name" in metrics: metrics[f"{event.kind}_name"] = metrics.pop("name") + if "kind" in metrics: metrics["backend_kind"] = metrics.pop("kind") + if "parent_id" in metrics: metrics["backend_parent_id"] = metrics.pop("parent_id") + self.trace.point("inference", event.kind, parent_id, **metrics) + + def _request_metadata(self, step: int) -> dict[str, Any]: + metadata = {"trace_id": self.trace.trace_id, "step": step, + "tool_round": sum(1 for event in self.events if event.kind == "tool_result"), + "resume_after_tool": bool(self.messages and self.messages[-1].get("role") == "tool")} + if self.metadata_provider is not None: metadata.update(dict(self.metadata_provider(self, step))) + return metadata + + def _generate(self, max_tokens: int, temperature: float) -> tuple[SessionEvent, ...]: + self.state = SessionState.GENERATING + self.partial_output = "" + self.pending_assistant_content = "" + self.last_finish_reason = None + started_at = len(self.events) + pieces: list[str] = [] + done_metrics: dict[str, Any] = {} + step = sum(1 for event in self.events if event.kind == "generation_started") + 1 + if step > self.max_agent_steps: raise RuntimeError("agent step limit exceeded") + with self.trace.span("agent", "model_turn", session_id=self.session_id, step=step) as parent: + self._emit("generation_started", backend=self.backend.name, step=step, capabilities=asdict(self.backend.capabilities)) + request = GenerationRequest(self.session_id, tuple(self.messages), tuple(self.tools.schemas()), max_tokens, temperature, + metadata=self._request_metadata(step)) + try: + for event in self.backend.stream(request): + if event.kind == "token": + pieces.append(event.text) + self.partial_output += event.text + self._emit("token", text=event.text, metrics=dict(event.metrics)) + token_metrics = dict(event.metrics) + if "name" in token_metrics: token_metrics["token_name"] = token_metrics.pop("name") + self.trace.duration("inference", "token", float(token_metrics.get("wall_ms", 0.0)), parent.event_id, + text=event.text, **token_metrics) + elif event.kind in {"prefill", "decode", "kernel", "metric", "hook"}: self._record_backend_event(event, parent.event_id) + elif event.kind == "tool_call" and event.tool_call is not None: + arguments = event.tool_call.get("arguments", {}) + if not isinstance(arguments, Mapping): raise ValueError("backend tool-call arguments must be an object") + self.pending_tool_call = ToolCall(str(event.tool_call.get("id") or uuid.uuid4().hex), + str(event.tool_call["name"]), dict(arguments)) + raw = event.tool_call.get("raw") + if isinstance(raw, str): self.pending_assistant_content = raw + self._emit("structured_tool_call", call=asdict(self.pending_tool_call), native=True) + elif event.kind == "done": + self.last_finish_reason = event.finish_reason + done_metrics = dict(event.metrics) + self._emit("generation_done", finish_reason=event.finish_reason, metrics=done_metrics) + except Exception as exc: + self.state = SessionState.FAILED + self._emit("error", error=str(exc), error_type=type(exc).__name__) + raise + output = "".join(pieces) + + if self.last_finish_reason == "cancelled": + discarded_tool_prefix = self._looks_like_partial_tool_protocol(output) + if output and not discarded_tool_prefix: self.messages.append({"role": "assistant", "content": output}) + self.pending_tool_call = None + self.pending_assistant_content = "" + self.partial_output = "" + self.state = SessionState.CANCELLED + self._emit("generation_cancelled", preserved_output=bool(output and not discarded_tool_prefix), + discarded_incomplete_tool_protocol=discarded_tool_prefix, generated_characters=len(output), + materialized_kv_tokens=done_metrics.get("materialized_kv_tokens"), cancel_stage=done_metrics.get("cancel_stage")) + return tuple(self.events[started_at:]) + + if self.pending_tool_call is None: + try: + self.pending_tool_call = parse_tool_call(output) + if self.pending_tool_call is not None: self.pending_assistant_content = output + except Exception as exc: self._emit("tool_parse_error", error=str(exc), raw=output) + if self.pending_tool_call is not None: + if not self.pending_assistant_content: self.pending_assistant_content = self._tool_content(self.pending_tool_call) + self.partial_output = "" + self.state = SessionState.AWAITING_TOOL_APPROVAL + self._emit("tool_approval_required", call=asdict(self.pending_tool_call), action=self.tools.spec(self.pending_tool_call.name).action.value) + else: + self.messages.append({"role": "assistant", "content": output}) + self.partial_output = "" + self.state = SessionState.COMPLETED + self._emit("message", role="assistant", content=output) + return tuple(self.events[started_at:]) + + def approve_tool(self, permission_token: str, max_tokens: int = 512) -> tuple[SessionEvent, ...]: + with self._lock: + if self.state is not SessionState.AWAITING_TOOL_APPROVAL or self.pending_tool_call is None: raise RuntimeError("no tool call awaits approval") + call = self.pending_tool_call + assistant_content = self.pending_assistant_content or self._tool_content(call) + started_at = len(self.events) + self.state = SessionState.RUNNING_TOOL + try: + with self.trace.span("tool", call.name, call_id=call.call_id) as span: + self._emit("tool_started", call=asdict(call)) + result = self.tools.execute(call, permission_token) + self.trace.point("tool", "tool_result", span.event_id, ok=result.ok, elapsed_ms=result.elapsed_ms) + except Exception as exc: + self.state = SessionState.AWAITING_TOOL_APPROVAL + self._emit("tool_authorization_failed", call_id=call.call_id, error=str(exc), error_type=type(exc).__name__) + raise + self._emit("tool_result", result=asdict(result)) + self.messages.append({"role": "assistant", "content": assistant_content}) + self.messages.append({"role": "tool", "name": call.name, "tool_call_id": call.call_id, "content": str(result.output)}) + self.pending_tool_call = None + self.pending_assistant_content = "" + if not result.ok: + self.state = SessionState.FAILED + return tuple(self.events[started_at:]) + self.state = SessionState.IDLE + self._generate(max_tokens, 0.0) + return tuple(self.events[started_at:]) + + def reject_tool(self, reason: str) -> SessionEvent: + with self._lock: + if self.pending_tool_call is None: raise RuntimeError("no pending tool call") + call = self.pending_tool_call + assistant_content = self.pending_assistant_content or self._tool_content(call) + self.messages.append({"role": "assistant", "content": assistant_content}) + self.messages.append({"role": "tool", "name": call.name, "tool_call_id": call.call_id, "content": f"User rejected tool call: {reason}"}) + self.pending_tool_call = None + self.pending_assistant_content = "" + self.state = SessionState.IDLE + return self._emit("tool_rejected", call_id=call.call_id, reason=reason) + + def events_after(self, sequence: int) -> list[dict[str, Any]]: + return [asdict(event) for event in self.events if event.sequence > sequence] + + def snapshot(self) -> dict[str, Any]: + return {"session_id": self.session_id, "state": self.state.value, "messages": list(self.messages), + "events": [asdict(x) for x in self.events], "pending_tool_call": asdict(self.pending_tool_call) if self.pending_tool_call else None, + "pending_assistant_content": self.pending_assistant_content, "partial_output": self.partial_output, + "last_finish_reason": self.last_finish_reason, "trace_id": self.trace.trace_id} diff --git a/extra/radeon_forge/runtime/stop_sequences.py b/extra/radeon_forge/runtime/stop_sequences.py new file mode 100644 index 0000000000000..41f6d26f234d9 --- /dev/null +++ b/extra/radeon_forge/runtime/stop_sequences.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Iterable + + +@dataclass(frozen=True) +class StopMatch: + text: str + matched: str | None = None + + +class StopSequenceMatcher: + """Incrementally hides stop strings even when they cross token boundaries.""" + def __init__(self, stops: Iterable[str]): + values = tuple(dict.fromkeys(str(stop) for stop in stops)) + if any(not stop for stop in values): raise ValueError("stop sequences must not be empty") + self.stops = values + self.buffer = "" + self.done = False + + def feed(self, text: str) -> StopMatch: + if self.done: raise RuntimeError("stop matcher is already complete") + self.buffer += text + earliest: tuple[int, str] | None = None + for stop in self.stops: + index = self.buffer.find(stop) + if index >= 0 and (earliest is None or index < earliest[0] or (index == earliest[0] and len(stop) > len(earliest[1]))): + earliest = (index, stop) + if earliest is not None: + index, stop = earliest + output = self.buffer[:index] + self.buffer = "" + self.done = True + return StopMatch(output, stop) + + keep = 0 + for stop in self.stops: + limit = min(len(stop) - 1, len(self.buffer)) + for length in range(limit, 0, -1): + if self.buffer.endswith(stop[:length]): + keep = max(keep, length) + break + if keep: + output, self.buffer = self.buffer[:-keep], self.buffer[-keep:] + else: + output, self.buffer = self.buffer, "" + return StopMatch(output) + + def finalize(self) -> str: + if self.done: return "" + output, self.buffer = self.buffer, "" + self.done = True + return output diff --git a/extra/radeon_forge/runtime/store.py b/extra/radeon_forge/runtime/store.py new file mode 100644 index 0000000000000..ceb6a6ec60f3c --- /dev/null +++ b/extra/radeon_forge/runtime/store.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +import json, os, tempfile, time +from dataclasses import asdict +from pathlib import Path +from typing import Any, Callable, Mapping + +from .backend import InferenceBackend +from .events import TraceEvent, TraceRecorder +from .session import AgentSession, SessionEvent, SessionState +from .tools import ToolCall, ToolRegistry + + +class SessionStore: + """Atomic, local-only persistence for agent sessions and unified traces.""" + FORMAT_VERSION = 1 + + def __init__(self, root: str | Path): + self.root = Path(root).resolve() + self.root.mkdir(parents=True, exist_ok=True) + + def _path(self, session_id: str) -> Path: + if not session_id or any(ch not in "0123456789abcdef" for ch in session_id.lower()): raise ValueError("invalid session id") + return self.root / f"{session_id}.json" + + @staticmethod + def _atomic_json(path: Path, payload: Mapping[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary = tempfile.mkstemp(prefix=path.name+".", suffix=".tmp", dir=path.parent) + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2, default=str) + handle.write("\n") + handle.flush(); os.fsync(handle.fileno()) + os.replace(temporary, path) + finally: + try: os.unlink(temporary) + except FileNotFoundError: pass + + def save(self, session: AgentSession) -> Path: + snapshot = session.snapshot() + payload = {"format_version":self.FORMAT_VERSION, "saved_at_s":time.time(), "session":snapshot, + "trace":session.trace.to_dict()} + path = self._path(session.session_id) + self._atomic_json(path, payload) + return path + + def delete(self, session_id: str) -> None: + try: self._path(session_id).unlink() + except FileNotFoundError: pass + + def list(self) -> list[dict[str, Any]]: + ret=[] + for path in sorted(self.root.glob("*.json"), key=lambda item:item.stat().st_mtime, reverse=True): + try: + payload=json.loads(path.read_text(encoding="utf-8")); session=payload["session"] + ret.append({"session_id":session["session_id"], "state":session.get("state","unknown"), + "saved_at_s":payload.get("saved_at_s"), "message_count":len(session.get("messages",[])), + "event_count":len(session.get("events",[])), "path":str(path)}) + except Exception: continue + return ret + + def load_payload(self, session_id: str) -> dict[str, Any]: + path=self._path(session_id) + payload=json.loads(path.read_text(encoding="utf-8")) + if int(payload.get("format_version",0)) != self.FORMAT_VERSION: raise ValueError("unsupported session checkpoint version") + return payload + + def restore(self, session_id: str, backend: InferenceBackend, tools: ToolRegistry, system_prompt: str, + metadata_provider: Callable[[AgentSession, int], Mapping[str, Any]] | None = None) -> AgentSession: + payload=self.load_payload(session_id); saved=payload["session"] + trace_payload=payload.get("trace", {}) + trace=TraceRecorder(str(trace_payload.get("trace_id") or saved.get("trace_id") or "")) + trace_events=[] + for item in trace_payload.get("events", []): + try: + trace_events.append(TraceEvent(str(item["event_id"]), str(item["trace_id"]), str(item["kind"]), str(item["name"]), + int(item["start_ns"]), int(item["end_ns"]) if item.get("end_ns") is not None else None, + str(item["parent_id"]) if item.get("parent_id") is not None else None, dict(item.get("attributes", {})))) + except Exception: continue + trace.extend(trace_events) + session=AgentSession(backend,tools,system_prompt,session_id=str(saved["session_id"]),trace=trace,metadata_provider=metadata_provider) + session.messages=[dict(item) for item in saved.get("messages", [])] + if not session.messages: session.messages=[{"role":"system","content":session._system_prompt()}] + session.events=[] + for item in saved.get("events", []): + try: session.events.append(SessionEvent(int(item["sequence"]),str(item["kind"]),dict(item.get("data",{})),float(item.get("timestamp_s",time.time())))) + except Exception: continue + session._sequence=max((item.sequence for item in session.events),default=0) + pending=saved.get("pending_tool_call") + session.pending_tool_call=(ToolCall(str(pending["call_id"]),str(pending["name"]),dict(pending.get("arguments",{}))) + if isinstance(pending,Mapping) else None) + session.pending_assistant_content=str(saved.get("pending_assistant_content", "")) + session.partial_output="" + session.last_finish_reason=saved.get("last_finish_reason") + prior_state=SessionState(str(saved.get("state",SessionState.IDLE.value))) + if prior_state in {SessionState.GENERATING,SessionState.RUNNING_TOOL}: + session.state=SessionState.IDLE + session._emit("session_recovered", prior_state=prior_state.value, action="incomplete operation was not replayed") + elif prior_state is SessionState.AWAITING_TOOL_APPROVAL and session.pending_tool_call is None: + session.state=SessionState.IDLE + session._emit("session_recovered", prior_state=prior_state.value, action="missing pending tool was cleared") + else: session.state=prior_state + return session diff --git a/extra/radeon_forge/runtime/tool_context.py b/extra/radeon_forge/runtime/tool_context.py new file mode 100644 index 0000000000000..dd9b98ea496c1 --- /dev/null +++ b/extra/radeon_forge/runtime/tool_context.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +import contextlib +from contextvars import ContextVar +from dataclasses import asdict, dataclass +from typing import Any, Iterator, Mapping + + +@dataclass(frozen=True) +class ToolExecutionContext: + session_id: str + trace_id: str + agent_step: int + tool_call_id: str + tool_name: str + attributes: Mapping[str, Any] + + +_CURRENT: ContextVar[ToolExecutionContext | None] = ContextVar("RADEON_FORGE_TOOL_CONTEXT", default=None) + + +@contextlib.contextmanager +def bind_tool_context(context: ToolExecutionContext) -> Iterator[None]: + token = _CURRENT.set(context) + try: yield + finally: _CURRENT.reset(token) + + +def current_tool_context(required: bool = False) -> ToolExecutionContext | None: + context = _CURRENT.get() + if required and context is None: raise RuntimeError("tool execution provenance is unavailable") + return context + + +def current_tool_context_dict() -> dict[str, Any]: + context = current_tool_context() + return asdict(context) if context is not None else {} diff --git a/extra/radeon_forge/runtime/tool_prompt.py b/extra/radeon_forge/runtime/tool_prompt.py new file mode 100644 index 0000000000000..59c0e762de695 --- /dev/null +++ b/extra/radeon_forge/runtime/tool_prompt.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +import json +from typing import Any, Mapping, Sequence + + +TOOL_PROTOCOL_MARKER = "Radeon Forge local tool protocol" + + +def render_tool_instruction(tools: Sequence[Mapping[str, Any]]) -> str: + rows = [] + for item in tools: + function = item.get("function", {}) if isinstance(item, Mapping) else {} + if not isinstance(function, Mapping) or not isinstance(function.get("name"), str): continue + rows.append({"name": function["name"], "description": str(function.get("description", "")), + "parameters": function.get("parameters", {"type":"object","properties":{}})}) + if not rows: return "" + return f"""{TOOL_PROTOCOL_MARKER}. +You may call only one of the tools listed below. When a tool is necessary, output exactly one object and no surrounding prose: +{{"name":"tool_name","arguments":{{...}}}} +Do not invent a result. Tool execution requires user approval and the result will be returned in a later message. +Available tools: +{json.dumps(rows, indent=2, sort_keys=True)}""" + + +def inject_tool_instruction(messages: Sequence[Mapping[str, Any]], tools: Sequence[Mapping[str, Any]]) -> list[dict[str, Any]]: + copied = [dict(message) for message in messages] + instruction = render_tool_instruction(tools) + if not instruction: return copied + if any(TOOL_PROTOCOL_MARKER in str(message.get("content", "")) for message in copied if message.get("role") == "system"): + return copied + index = next((i for i, message in enumerate(copied) if message.get("role") == "system"), None) + if index is None: copied.insert(0, {"role":"system", "content":instruction}) + else: copied[index]["content"] = str(copied[index].get("content", "")).rstrip() + "\n\n" + instruction + return copied diff --git a/extra/radeon_forge/runtime/tool_stream.py b/extra/radeon_forge/runtime/tool_stream.py new file mode 100644 index 0000000000000..f90effef378c2 --- /dev/null +++ b/extra/radeon_forge/runtime/tool_stream.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +import json, re, uuid +from dataclasses import dataclass +from typing import Iterable, Mapping, Any + + +_TOOL_CALL = re.compile(r"\s*\s*(\{.*\})\s*\s*", re.S) + + +class ToolProtocolError(ValueError): pass + + +@dataclass(frozen=True) +class ParsedToolCall: + call_id: str + name: str + arguments: Mapping[str, Any] + raw: str + + def to_event(self) -> dict[str, Any]: + return {"id": self.call_id, "name": self.name, "arguments": dict(self.arguments), "raw": self.raw} + + +class ToolStreamParser: + """Incrementally detects the exact local tool-call protocol. + + It is not a grammar decoder yet; it is a bounded early-stop and validation + layer. Malformed or unknown tool calls fail closed rather than being executed + as best-effort text. + """ + def __init__(self, allowed_names: Iterable[str], max_bytes: int = 65536): + self.allowed_names = frozenset(str(name) for name in allowed_names) + self.max_bytes = max_bytes + self.buffer = "" + self.completed = False + + def feed(self, text: str) -> ParsedToolCall | None: + if self.completed: raise ToolProtocolError("tool stream is already complete") + self.buffer += text + if len(self.buffer.encode("utf-8")) > self.max_bytes: raise ToolProtocolError("tool-call stream exceeds size limit") + if "" not in self.buffer: return None + match = _TOOL_CALL.fullmatch(self.buffer) + if match is None: raise ToolProtocolError("malformed tool call: output must contain exactly one wrapped JSON object") + try: payload = json.loads(match.group(1)) + except json.JSONDecodeError as exc: raise ToolProtocolError(f"tool call contains invalid JSON: {exc}") from exc + if not isinstance(payload, dict): raise ToolProtocolError("tool-call payload must be an object") + name, arguments = payload.get("name"), payload.get("arguments", {}) + if not isinstance(name, str) or not name: raise ToolProtocolError("tool call requires a non-empty string name") + if name not in self.allowed_names: raise ToolProtocolError(f"unknown or unavailable local tool {name!r}") + if not isinstance(arguments, dict): raise ToolProtocolError("tool call arguments must be an object") + call_id = payload.get("id") or uuid.uuid4().hex + if not isinstance(call_id, str): raise ToolProtocolError("tool call id must be a string") + self.completed = True + return ParsedToolCall(call_id, name, arguments, self.buffer) + + def finalize(self) -> None: + stripped = self.buffer.lstrip() + if stripped.startswith("") and not self.completed: + raise ToolProtocolError("generation ended with an incomplete tool call") diff --git a/extra/radeon_forge/runtime/tools.py b/extra/radeon_forge/runtime/tools.py new file mode 100644 index 0000000000000..0ef7d182f1d25 --- /dev/null +++ b/extra/radeon_forge/runtime/tools.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import json, os, re, shlex, subprocess, time, uuid +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Mapping, Sequence + +from ..permissions import Action, PermissionController + + +@dataclass(frozen=True) +class ToolSpec: + name: str + description: str + parameters: Mapping[str, Any] + action: Action = Action.INSPECT + + def openai_schema(self) -> dict[str, Any]: + return {"type": "function", "function": {"name": self.name, "description": self.description, "parameters": dict(self.parameters)}} + + +@dataclass(frozen=True) +class ToolCall: + call_id: str + name: str + arguments: Mapping[str, Any] + + +@dataclass(frozen=True) +class ToolResult: + call_id: str + name: str + ok: bool + output: Any + elapsed_ms: float + + +@dataclass +class _RegisteredTool: + spec: ToolSpec + fn: Callable[[Mapping[str, Any]], Any] + + +class ToolRegistry: + def __init__(self, permissions: PermissionController): + self.permissions = permissions + self._tools: dict[str, _RegisteredTool] = {} + + def register(self, spec: ToolSpec, fn: Callable[[Mapping[str, Any]], Any]) -> None: + if spec.name in self._tools: raise ValueError(f"duplicate tool {spec.name}") + self._tools[spec.name] = _RegisteredTool(spec, fn) + + def schemas(self) -> list[dict[str, Any]]: return [tool.spec.openai_schema() for tool in self._tools.values()] + def spec(self, name: str) -> ToolSpec: return self._tools[name].spec + + def execute(self, call: ToolCall, permission_token: str | None) -> ToolResult: + if call.name not in self._tools: return ToolResult(call.call_id, call.name, False, {"error": "unknown tool"}, 0.0) + tool = self._tools[call.name] + self.permissions.authorize(permission_token, tool.spec.action) + started = time.perf_counter_ns() + try: output, ok = tool.fn(call.arguments), True + except Exception as exc: output, ok = {"error": str(exc), "type": type(exc).__name__}, False + return ToolResult(call.call_id, call.name, ok, output, (time.perf_counter_ns() - started) / 1e6) + + +class WorkspaceTools: + """Safe-by-construction tools rooted in one private repository workspace.""" + def __init__(self, root: str | Path, command_allowlist: Sequence[str] = ("python", "python3", "pytest", "git", "rg")): + self.root = Path(root).resolve() + self.command_allowlist = frozenset(command_allowlist) + + def _path(self, value: str) -> Path: + path = (self.root / value).resolve() + if path != self.root and self.root not in path.parents: raise ValueError("path escapes workspace") + return path + + def list_files(self, args: Mapping[str, Any]) -> Any: + base = self._path(str(args.get("path", "."))) + limit = max(1, min(int(args.get("limit", 200)), 2000)) + files = [] + for path in base.rglob("*"): + if path.is_file() and ".git" not in path.parts: + files.append(str(path.relative_to(self.root))) + if len(files) >= limit: break + return {"files": files, "truncated": len(files) >= limit} + + def read_file(self, args: Mapping[str, Any]) -> Any: + path = self._path(str(args["path"])) + start, end = max(1, int(args.get("start_line", 1))), int(args.get("end_line", 400)) + if end < start or end - start > 2000: raise ValueError("invalid line range") + lines = path.read_text(encoding="utf-8", errors="replace").splitlines() + return {"path": str(path.relative_to(self.root)), "start_line": start, "end_line": min(end, len(lines)), + "content": "\n".join(f"{i}: {lines[i-1]}" for i in range(start, min(end, len(lines)) + 1))} + + def search_text(self, args: Mapping[str, Any]) -> Any: + pattern = re.compile(str(args["query"])) + base = self._path(str(args.get("path", "."))) + limit = max(1, min(int(args.get("limit", 50)), 500)) + hits = [] + for path in base.rglob("*"): + if not path.is_file() or ".git" in path.parts or path.stat().st_size > 2_000_000: continue + try: lines = path.read_text(encoding="utf-8", errors="replace").splitlines() + except OSError: continue + for number, line in enumerate(lines, 1): + if pattern.search(line): + hits.append({"path": str(path.relative_to(self.root)), "line": number, "text": line[:500]}) + if len(hits) >= limit: return {"hits": hits, "truncated": True} + return {"hits": hits, "truncated": False} + + def run_command(self, args: Mapping[str, Any]) -> Any: + command = args.get("command") + argv = shlex.split(command) if isinstance(command, str) else [str(x) for x in command] + if not argv or Path(argv[0]).name not in self.command_allowlist: raise ValueError("command is not allowlisted") + timeout = max(1, min(int(args.get("timeout_seconds", 60)), 600)) + proc = subprocess.run(argv, cwd=self.root, text=True, capture_output=True, timeout=timeout, + env={**os.environ, "PYTHONUNBUFFERED": "1"}) + return {"argv": argv, "returncode": proc.returncode, "stdout": proc.stdout[-20000:], "stderr": proc.stderr[-20000:]} + + def install(self, registry: ToolRegistry) -> None: + registry.register(ToolSpec("list_files", "List files inside the private workspace", {"type":"object","properties":{"path":{"type":"string"},"limit":{"type":"integer"}}}), self.list_files) + registry.register(ToolSpec("read_file", "Read a line range from a workspace file", {"type":"object","required":["path"],"properties":{"path":{"type":"string"},"start_line":{"type":"integer"},"end_line":{"type":"integer"}}}), self.read_file) + registry.register(ToolSpec("search_text", "Regex-search private workspace text", {"type":"object","required":["query"],"properties":{"query":{"type":"string"},"path":{"type":"string"},"limit":{"type":"integer"}}}), self.search_text) + registry.register(ToolSpec("run_command", "Run an allowlisted local development command", {"type":"object","required":["command"],"properties":{"command":{"type":["string","array"]},"timeout_seconds":{"type":"integer"}}}, Action.BENCHMARK), self.run_command) + + +def parse_tool_call(text: str) -> ToolCall | None: + match = re.fullmatch(r"\s*\s*(\{.*\})\s*\s*", text, flags=re.S) + if not match: return None + payload = json.loads(match.group(1)) + name = payload.get("name") + arguments = payload.get("arguments", {}) + if not isinstance(name, str) or not isinstance(arguments, dict): raise ValueError("invalid tool call payload") + return ToolCall(str(payload.get("id") or uuid.uuid4().hex), name, arguments) diff --git a/extra/radeon_forge/synthesis/__init__.py b/extra/radeon_forge/synthesis/__init__.py new file mode 100644 index 0000000000000..314e1a20dc344 --- /dev/null +++ b/extra/radeon_forge/synthesis/__init__.py @@ -0,0 +1,15 @@ +from .autotune import SearchPlan, run_autotune +from .defaults import install_default_specs +from .deployment import RUNTIME_ADAPTERS, STATEFUL_LAYERS, SafeHookRegistry +from .hooks import (ActiveHook, CompatibilityReport, ExecutionContext, ExecutionStage, HookDescriptor, HookLayer, HookMode, + HookRegistry, RuntimeFingerprint, StagePredicate, check_compatibility) +from .portable import export_recipe_with_hook +from .recipe import ForgeRecipe, InstalledRecipe, RecipeArtifact, RecipeLibrary, export_recipe +from .tools import OptimizationTools +from .workspace import CandidateRecord, CandidateWorkspace, KernelSpec + +__all__ = ["ActiveHook", "CandidateRecord", "CandidateWorkspace", "CompatibilityReport", "ExecutionContext", "ExecutionStage", + "ForgeRecipe", "HookDescriptor", "HookLayer", "HookMode", "HookRegistry", "InstalledRecipe", "KernelSpec", + "OptimizationTools", "RUNTIME_ADAPTERS", "RecipeArtifact", "RecipeLibrary", "RuntimeFingerprint", "STATEFUL_LAYERS", + "SafeHookRegistry", "SearchPlan", "StagePredicate", "check_compatibility", "export_recipe", "export_recipe_with_hook", + "install_default_specs", "run_autotune"] diff --git a/extra/radeon_forge/synthesis/autotune.py b/extra/radeon_forge/synthesis/autotune.py new file mode 100644 index 0000000000000..a0c979188fc3b --- /dev/null +++ b/extra/radeon_forge/synthesis/autotune.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +import hashlib, itertools, json, os +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Mapping, Sequence + +from ..command_backend import CommandHarness +from ..contracts import Candidate, Objective, WorkloadContract +from ..ledger import ExperimentLedger +from ..tuner import TuningSummary, successive_halving +from .hooks import HookDescriptor +from .workspace import CandidateWorkspace + + +@dataclass(frozen=True) +class SearchPlan: + axes: Mapping[str, tuple[int | float | str | bool, ...]] + budgets: tuple[int, ...] = (3, 10, 30) + reduction: int = 3 + max_candidates: int = 128 + timeout_seconds: int = 900 + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> SearchPlan: + raw = dict(value) + axes_value = raw.get("axes", {}) + if not isinstance(axes_value, Mapping) or not axes_value: raise ValueError("autotune search requires a non-empty axes table") + axes: dict[str, tuple[int | float | str | bool, ...]] = {} + for name, values in axes_value.items(): + if not isinstance(values, Sequence) or isinstance(values, (str, bytes, bytearray)) or not values: + raise ValueError(f"search axis {name!r} must be a non-empty array") + converted = tuple(values) + if not all(isinstance(item, (int, float, str, bool)) for item in converted): + raise ValueError(f"search axis {name!r} contains an unsupported value") + axes[str(name)] = converted + budgets = tuple(int(x) for x in raw.get("budgets", (3, 10, 30))) + if not budgets or any(x <= 0 for x in budgets): raise ValueError("autotune budgets must be positive") + reduction = int(raw.get("reduction", 3)) + max_candidates = int(raw.get("max_candidates", 128)) + timeout_seconds = int(raw.get("timeout_seconds", 900)) + if reduction < 2 or max_candidates < 1 or timeout_seconds < 1: raise ValueError("invalid autotune search limits") + plan = cls(axes, budgets, reduction, max_candidates, timeout_seconds) + if plan.cardinality > max_candidates: raise ValueError(f"search has {plan.cardinality} candidates, maximum is {max_candidates}") + return plan + + @property + def cardinality(self) -> int: + total = 1 + for values in self.axes.values(): total *= len(values) + return total + + def parameters(self) -> list[dict[str, int | float | str | bool]]: + names = tuple(self.axes) + return [dict(zip(names, values)) for values in itertools.product(*(self.axes[name] for name in names))] + + +def _candidate_id(base_id: str, parameters: Mapping[str, Any]) -> str: + digest = hashlib.sha256(json.dumps(parameters, sort_keys=True, default=str).encode()).hexdigest()[:12] + return f"{base_id}-tune-{digest}" + + +def _contract(spec) -> WorkloadContract: + acceptance = spec.metadata.get("recipe_acceptance", spec.metadata.get("acceptance", {})) + if not isinstance(acceptance, Mapping): acceptance = {} + objective_text = str(spec.metadata.get("search_objective", "p95_latency_us")) + try: objective = Objective(objective_text) + except ValueError: objective = Objective.P95_LATENCY_US + return WorkloadContract( + name=f"{spec.name}:{HookDescriptor.from_spec(spec).when.signature}", target=spec.target, objective=objective, + max_abs_error=float(acceptance.get("max_abs_error", acceptance.get("maximum_error", 1e-3))), + max_rel_error=float(acceptance.get("max_rel_error", 1e-3)), + max_vram_bytes=int(acceptance["max_vram_bytes"]) if acceptance.get("max_vram_bytes") is not None else None, + max_vgprs=int(acceptance["max_vgprs"]) if acceptance.get("max_vgprs") is not None else None, + max_lds_bytes=int(acceptance["max_lds_bytes"]) if acceptance.get("max_lds_bytes") is not None else None, + forbid_spills=bool(acceptance.get("forbid_spills", True)), + metadata={"hook": asdict(HookDescriptor.from_spec(spec)), "objective_text": spec.objective}, + ) + + +def run_autotune(workspace: CandidateWorkspace, candidate_id: str, project_root: str | Path, + plan: SearchPlan, ledger_path: str | Path | None = None) -> tuple[Any, TuningSummary]: + """Tune one structural implementation without allowing speed to bypass correctness.""" + record = workspace.load_candidate(candidate_id) + if record.status != "mockgpu_passed": raise ValueError("candidate must pass MockGPU before stage-specific W7900 autotuning") + spec = workspace.load_spec(record.spec_id) + if not spec.hardware_command: raise ValueError("spec has no hardware command for autotuning") + root = Path(project_root).resolve() + command = workspace.render_command(spec.hardware_command, spec, record, root) + cwd = str(spec.metadata.get("recipe_bundle", root)) + environment = {"DEV": "AMD", "RADEON_FORGE_CANDIDATE": record.source_path, + "RADEON_FORGE_RECIPE_BUNDLE": str(spec.metadata.get("recipe_bundle", "")), + "RADEON_FORGE_EXECUTION_STAGES": ",".join(x.value for x in HookDescriptor.from_spec(spec).when.stages), + "PYTHONUNBUFFERED": "1"} + harness = CommandHarness(command, cwd=cwd, env={**os.environ, **environment}, timeout_seconds=plan.timeout_seconds) + candidates = [Candidate(_candidate_id(record.candidate_id, parameters), spec.name, parameters, + source_path=record.source_path, hypothesis=record.hypothesis, parent_id=record.candidate_id) + for parameters in plan.parameters()] + ledger = ExperimentLedger(ledger_path or (workspace.root / "autotune" / f"{record.candidate_id}.jsonl")) + summary = successive_halving(candidates, harness, _contract(spec), plan.budgets, plan.reduction, ledger) + evidence = {"search_plan": asdict(plan), "rounds": summary.rounds, "evaluated_trials": len(summary.evaluated), + "ledger": str(ledger.path), "execution_hook": asdict(HookDescriptor.from_spec(spec)), + "winner": summary.winner.to_dict() if summary.winner is not None else None, + "rejections": [{"candidate_id": result.candidate.candidate_id, "parameters": dict(result.candidate.parameters), + "reason": result.rejected_reason or result.correctness.reason, + "feasible": result.is_feasible(_contract(spec))} for result in summary.evaluated]} + if summary.winner is None: + updated = workspace.update(record.candidate_id, "autotune_failed", {"autotune": evidence}) + else: + updated = workspace.update(record.candidate_id, "hardware_passed", {"autotune": evidence, + "selected_parameters": dict(summary.winner.candidate.parameters), "hardware": summary.winner.to_dict()}) + return updated, summary diff --git a/extra/radeon_forge/synthesis/defaults.py b/extra/radeon_forge/synthesis/defaults.py new file mode 100644 index 0000000000000..0ff82351288c4 --- /dev/null +++ b/extra/radeon_forge/synthesis/defaults.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from .workspace import CandidateWorkspace, KernelSpec + + +def install_default_specs(workspace: CandidateWorkspace) -> list[KernelSpec]: + specs = [ + KernelSpec( + name="rdna3-vopd-gemm-schedule", + operation="C = A @ B; batch-one projection/GEMM scheduling research", + target="gfx1100", language="python", extension=".py", entrypoint="build_kernel", + shapes={"M":"multiple of 128", "N":"multiple of 128", "K":"multiple of 128"}, + dtypes={"A":"float32", "B":"float32", "C":"float32"}, + invariants=("Numerically match tinygrad matmul within MSE 1e-6", "No unsupported RDNA3 instructions", + "No out-of-workspace file access", "MockGPU timing is never used as a speed metric"), + objective="minimize median and P95 W7900 latency without correctness loss", + mockgpu_command=("python3", "{candidate}"), hardware_command=("python3", "{candidate}"), + metadata={"seed":"extra/gemm/amd_asm_matmul.py", "role":"infrastructure and schedule-synthesis gate", + "hook":{"layer":"kernel", "target":"batch1_projection_gemm", "mode":"replace", "adapter":"request_metadata", + "selector":{"program_name":"generated_gemm"}, + "when":{"stages":["prefill", "first_token", "decode"], "batch_sizes":[1]}, + "exclusive_group":"projection_gemm", + "description":"Validated kernel selection forwarded to a backend-specific kernel adapter."}}, + ), + KernelSpec( + name="batch1-decode-megakernel", + operation="One specialized transformer decode block for batch-one private-agent inference", + target="gfx1100", language="python", extension=".py", entrypoint="build_kernel", + shapes={"batch":1, "sequence":"decode token", "model":"fixed by workload contract"}, + dtypes={"activations":"bf16/fp16", "accumulation":"fp32 where required"}, + invariants=("Match the trusted tinygrad block reference", "Preserve KV-cache semantics", "No spills unless explicitly allowed", + "Generated code is disposable; specification and oracle are authoritative"), + objective="minimize inter-token latency and kernel launches/token", + metadata={"status":"oracle command must be bound after selecting the model", "role":"final technical target", + "hook":{"layer":"transformer_block", "target":"llama.decode.block", "mode":"replace", + "adapter":"python_transformer_block", "selector":{"indices":"all"}, + "when":{"stages":["decode"], "batch_sizes":[1], "min_generated_token_index":1}, + "exclusive_group":"decode_transformer_block", + "description":"Replace selected tinygrad Llama blocks only during steady batch-one decode."}}, + ), + ] + existing = {x.spec_id for x in workspace.specs()} + for spec in specs: + if spec.spec_id not in existing: workspace.save_spec(spec) + return specs diff --git a/extra/radeon_forge/synthesis/deployment.py b/extra/radeon_forge/synthesis/deployment.py new file mode 100644 index 0000000000000..c1b54da587463 --- /dev/null +++ b/extra/radeon_forge/synthesis/deployment.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from .hooks import (ActiveHook, HookActivationError, HookDescriptor, HookLayer, HookRegistry, RuntimeFingerprint) +from .workspace import CandidateWorkspace + + +STATEFUL_LAYERS = frozenset({HookLayer.MODEL, HookLayer.TRANSFORMER_BLOCK, HookLayer.KV_CACHE}) +RUNTIME_ADAPTERS = frozenset({"python_transformer_block"}) + + +class SafeHookRegistry(HookRegistry): + """Deployment policy layered over the generic stage resolver. + + The generic registry is useful for tests and future adapters. The engine uses + this stricter registry so metadata-only or unimplemented adapters cannot be + presented as active optimizations, and stateful phase changes cannot deploy + without a held-out transition oracle. + """ + SUPPORTED_ADAPTERS = RUNTIME_ADAPTERS + + def __init__(self, workspace: CandidateWorkspace): super().__init__(workspace) + + def activate(self, candidate_id: str, fingerprint: RuntimeFingerprint, reason: str) -> ActiveHook: + candidate = self.workspace.load_candidate(candidate_id) + spec = self.workspace.load_spec(candidate.spec_id) + descriptor = HookDescriptor.from_spec(spec) + if descriptor.adapter not in self.SUPPORTED_ADAPTERS: + raise HookActivationError( + f"runtime adapter {descriptor.adapter!r} is not executable in this engine build; " + "the recipe remains usable as profiling/search knowledge" + ) + if descriptor.layer in STATEFUL_LAYERS: + heldout = tuple(spec.metadata.get("heldout_command", ())) + if not heldout: + raise HookActivationError( + f"{descriptor.layer.value} hooks require a held-out phase-transition oracle before deployment " + "(for example prefill -> first_token -> decode and tool-resume continuity)" + ) + if candidate.status not in {"heldout_passed", "accepted"}: + raise HookActivationError( + f"stateful hook candidate status {candidate.status!r} has not passed its phase-transition oracle" + ) + return super().activate(candidate_id, fingerprint, reason) diff --git a/extra/radeon_forge/synthesis/hooks.py b/extra/radeon_forge/synthesis/hooks.py new file mode 100644 index 0000000000000..f98253d83db7b --- /dev/null +++ b/extra/radeon_forge/synthesis/hooks.py @@ -0,0 +1,361 @@ +from __future__ import annotations + +import hashlib, json, threading, time, uuid +from dataclasses import asdict, dataclass, field +from enum import Enum +from pathlib import Path +from typing import Any, Mapping, Sequence + +from .workspace import CandidateWorkspace, KernelSpec + + +class HookLayer(str, Enum): + AGENT_LOOP = "agent_loop" + SCHEDULER = "scheduler" + MODEL = "model" + TRANSFORMER_BLOCK = "transformer_block" + SUBGRAPH = "subgraph" + KERNEL = "kernel" + KV_CACHE = "kv_cache" + SAMPLER = "sampler" + + +class HookMode(str, Enum): + BEFORE = "before" + AFTER = "after" + REPLACE = "replace" + WRAP = "wrap" + + +class ExecutionStage(str, Enum): + """Coarse states in one agent/inference execution. + + Stage is deliberately orthogonal to HookLayer. A kernel or block replacement + can be valid for decode but harmful for prefill; a scheduler optimization may + only apply while resuming after a tool call; a KV optimization may only apply + while appending to a warm cache. + """ + ANY = "any" + SESSION_START = "session_start" + PREFILL = "prefill" + FIRST_TOKEN = "first_token" + DECODE = "decode" + TOOL_EXECUTION = "tool_execution" + TOOL_RESUME = "tool_resume" + KV_APPEND = "kv_append" + SAMPLING = "sampling" + SESSION_END = "session_end" + + +@dataclass(frozen=True) +class ExecutionContext: + stage: ExecutionStage + batch_size: int = 1 + prompt_tokens: int = 0 + context_tokens: int = 0 + generated_token_index: int = -1 + prefix_reused_tokens: int = 0 + tool_round: int = 0 + warm: bool = False + attributes: Mapping[str, Any] = field(default_factory=dict) + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> ExecutionContext: + raw = dict(value) + stage = ExecutionStage(str(raw.pop("stage"))) + known = {key: raw.pop(key, default) for key, default in ( + ("batch_size", 1), ("prompt_tokens", 0), ("context_tokens", 0), ("generated_token_index", -1), + ("prefix_reused_tokens", 0), ("tool_round", 0), ("warm", False))} + return cls(stage, int(known["batch_size"]), int(known["prompt_tokens"]), int(known["context_tokens"]), + int(known["generated_token_index"]), int(known["prefix_reused_tokens"]), int(known["tool_round"]), + bool(known["warm"]), raw) + + def flattened(self) -> dict[str, Any]: + return {"stage": self.stage.value, "batch_size": self.batch_size, "prompt_tokens": self.prompt_tokens, + "context_tokens": self.context_tokens, "generated_token_index": self.generated_token_index, + "prefix_reused_tokens": self.prefix_reused_tokens, "tool_round": self.tool_round, "warm": self.warm, + "prefix_cache": "hit" if self.prefix_reused_tokens > 0 else "miss", **dict(self.attributes)} + + +@dataclass(frozen=True) +class StagePredicate: + stages: tuple[ExecutionStage, ...] = (ExecutionStage.ANY,) + min_context_tokens: int | None = None + max_context_tokens: int | None = None + min_prompt_tokens: int | None = None + max_prompt_tokens: int | None = None + min_generated_token_index: int | None = None + max_generated_token_index: int | None = None + batch_sizes: tuple[int, ...] = () + prefix_cache: str = "any" # any, hit, miss + warm: bool | None = None + conditions: Mapping[str, Any] = field(default_factory=dict) + + @classmethod + def from_mapping(cls, value: Mapping[str, Any] | None) -> StagePredicate: + raw = dict(value or {}) + stage_value = raw.pop("stages", raw.pop("stage", ["any"])) + if isinstance(stage_value, str): stage_value = [stage_value] + if not isinstance(stage_value, Sequence): raise ValueError("hook stages must be a string or array") + stages = tuple(ExecutionStage(str(x)) for x in stage_value) + batch_value = raw.pop("batch_sizes", raw.pop("batch_size", ())) + if isinstance(batch_value, int): batch_value = [batch_value] + if not isinstance(batch_value, Sequence): raise ValueError("hook batch_sizes must be an integer or array") + prefix_cache = str(raw.pop("prefix_cache", "any")) + if prefix_cache not in {"any", "hit", "miss"}: raise ValueError("prefix_cache must be any, hit, or miss") + warm = raw.pop("warm", None) + if warm is not None and not isinstance(warm, bool): raise ValueError("warm must be boolean") + fields = {} + for name in ("min_context_tokens", "max_context_tokens", "min_prompt_tokens", "max_prompt_tokens", + "min_generated_token_index", "max_generated_token_index"): + item = raw.pop(name, None) + fields[name] = None if item is None else int(item) + explicit_conditions = raw.pop("conditions", {}) + if not isinstance(explicit_conditions, Mapping): raise ValueError("hook conditions must be a table") + conditions = {**raw, **dict(explicit_conditions)} + return cls(stages, **fields, batch_sizes=tuple(int(x) for x in batch_value), prefix_cache=prefix_cache, warm=warm, + conditions=conditions) + + @property + def signature(self) -> str: + return hashlib.sha256(json.dumps(asdict(self), sort_keys=True, default=str).encode()).hexdigest()[:12] + + def matches(self, context: ExecutionContext) -> bool: + if ExecutionStage.ANY not in self.stages and context.stage not in self.stages: return False + if self.min_context_tokens is not None and context.context_tokens < self.min_context_tokens: return False + if self.max_context_tokens is not None and context.context_tokens > self.max_context_tokens: return False + if self.min_prompt_tokens is not None and context.prompt_tokens < self.min_prompt_tokens: return False + if self.max_prompt_tokens is not None and context.prompt_tokens > self.max_prompt_tokens: return False + if self.min_generated_token_index is not None and context.generated_token_index < self.min_generated_token_index: return False + if self.max_generated_token_index is not None and context.generated_token_index > self.max_generated_token_index: return False + if self.batch_sizes and context.batch_size not in self.batch_sizes: return False + if self.prefix_cache != "any" and context.flattened()["prefix_cache"] != self.prefix_cache: return False + if self.warm is not None and context.warm is not self.warm: return False + flat = context.flattened() + return all(_matches(expected, flat.get(str(key))) for key, expected in self.conditions.items()) + + +@dataclass(frozen=True) +class HookDescriptor: + layer: HookLayer + target: str + mode: HookMode = HookMode.REPLACE + adapter: str = "request_metadata" + selector: Mapping[str, Any] = field(default_factory=dict) + when: StagePredicate = field(default_factory=StagePredicate) + priority: int = 0 + exclusive_group: str = "" + description: str = "" + + @classmethod + def from_spec(cls, spec: KernelSpec) -> HookDescriptor: + raw = spec.metadata.get("hook", {}) + if not isinstance(raw, Mapping): raise ValueError("spec metadata hook must be a table") + layer = HookLayer(str(raw.get("layer", "kernel"))) + target = str(raw.get("target", spec.operation)).strip() + if not target: raise ValueError("hook target must not be empty") + selector = raw.get("selector", {}) + if not isinstance(selector, Mapping): raise ValueError("hook selector must be a table") + when_value = raw.get("when", {key: raw[key] for key in ("stage", "stages") if key in raw}) + predicate = StagePredicate.from_mapping(when_value if isinstance(when_value, Mapping) else {"stages": when_value}) + group = str(raw.get("exclusive_group", f"{layer.value}:{target}")) + return cls(layer, target, HookMode(str(raw.get("mode", "replace"))), str(raw.get("adapter", "request_metadata")), + dict(selector), predicate, int(raw.get("priority", 0)), group, str(raw.get("description", ""))) + + @property + def slot_key(self) -> str: return f"{self.exclusive_group}:{self.when.signature}" + + +@dataclass(frozen=True) +class RuntimeFingerprint: + architecture: str = "" + gpu: str = "" + runtime: str = "" + runtime_revision: str = "" + model_family: str = "" + model_hash: str = "" + dtype: str = "" + shapes: Mapping[str, Any] = field(default_factory=dict) + extra: Mapping[str, Any] = field(default_factory=dict) + + @classmethod + def from_mapping(cls, value: Mapping[str, Any] | None) -> RuntimeFingerprint: + raw = dict(value or {}) + known = {key: raw.pop(key, "") for key in ("architecture", "gpu", "runtime", "runtime_revision", "model_family", "model_hash", "dtype")} + shapes = raw.pop("shapes", {}) + return cls(**known, shapes=dict(shapes) if isinstance(shapes, Mapping) else {}, extra=raw) + + def flattened(self) -> dict[str, Any]: + ret = {"architecture": self.architecture, "gpu": self.gpu, "runtime": self.runtime, + "runtime_revision": self.runtime_revision, "model_family": self.model_family, + "model_hash": self.model_hash, "dtype": self.dtype} + ret.update({f"shape.{key}": value for key, value in self.shapes.items()}) + ret.update(self.extra) + return ret + + +@dataclass(frozen=True) +class CompatibilityReport: + compatible: bool + exact: tuple[str, ...] + unknown: tuple[str, ...] + mismatches: tuple[str, ...] + + +def _matches(expected: Any, actual: Any) -> bool: + if expected in (None, "", "*"): return True + if isinstance(expected, Sequence) and not isinstance(expected, (str, bytes, bytearray)): return actual in expected + if isinstance(expected, str) and expected.endswith("*"): return str(actual).startswith(expected[:-1]) + return expected == actual + + +def check_compatibility(spec: KernelSpec, fingerprint: RuntimeFingerprint) -> CompatibilityReport: + expected = spec.metadata.get("recipe_compatibility", spec.metadata.get("compatibility", {})) + if not isinstance(expected, Mapping): raise ValueError("compatibility contract must be a table") + actual = fingerprint.flattened() + exact, unknown, mismatches = [], [], [] + if spec.target: + arch = actual.get("architecture", "") + if not arch: unknown.append("architecture") + elif _matches(spec.target, arch): exact.append("architecture") + else: mismatches.append(f"architecture expected={spec.target!r} actual={arch!r}") + for key, value in expected.items(): + key = str(key) + current = actual.get(key) + if current in (None, ""): unknown.append(key) + elif _matches(value, current): exact.append(key) + else: mismatches.append(f"{key} expected={value!r} actual={current!r}") + return CompatibilityReport(not mismatches, tuple(exact), tuple(unknown), tuple(mismatches)) + + +@dataclass(frozen=True) +class ActiveHook: + activation_id: str + candidate_id: str + spec_id: str + source_path: str + source_sha256: str + descriptor: HookDescriptor + compatibility: CompatibilityReport + activated_at_s: float + reason: str + previous_activation_id: str | None = None + + +class HookActivationError(RuntimeError): pass + + +def _descriptor_from_dict(desc: Mapping[str, Any]) -> HookDescriptor: + when = desc.get("when", {}) + return HookDescriptor(HookLayer(desc["layer"]), str(desc["target"]), HookMode(desc["mode"]), str(desc["adapter"]), + dict(desc.get("selector", {})), StagePredicate.from_mapping(when), int(desc.get("priority", 0)), + str(desc.get("exclusive_group", "")), str(desc.get("description", ""))) + + +class HookRegistry: + """Persistent, rollback-safe, execution-state-aware optimization registry.""" + SUPPORTED_ADAPTERS = frozenset({"request_metadata", "python_transformer_block"}) + + def __init__(self, workspace: CandidateWorkspace): + self.workspace = workspace + self.root = workspace.root / "hooks" + self.root.mkdir(parents=True, exist_ok=True) + self.active_path, self.history_path = self.root / "active.json", self.root / "history.jsonl" + self._lock = threading.RLock() + if not self.active_path.exists(): self.active_path.write_text("[]\n", encoding="utf-8") + + def _load_active(self) -> list[ActiveHook]: + try: payload = json.loads(self.active_path.read_text(encoding="utf-8")) + except Exception: payload = [] + ret = [] + for item in payload: + try: + comp = item["compatibility"] + ret.append(ActiveHook(item["activation_id"], item["candidate_id"], item["spec_id"], item["source_path"], + item["source_sha256"], _descriptor_from_dict(item["descriptor"]), + CompatibilityReport(bool(comp["compatible"]), tuple(comp.get("exact", ())), + tuple(comp.get("unknown", ())), tuple(comp.get("mismatches", ()))), float(item["activated_at_s"]), + item["reason"], item.get("previous_activation_id"))) + except Exception: continue + return ret + + def _save_active(self, hooks: Sequence[ActiveHook]) -> None: + temp = self.active_path.with_suffix(".tmp") + temp.write_text(json.dumps([asdict(x) for x in hooks], indent=2, default=str) + "\n", encoding="utf-8") + temp.replace(self.active_path) + + def _history(self, event: str, **data: Any) -> None: + with self.history_path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps({"event": event, "timestamp_s": time.time(), **data}, default=str) + "\n") + + def active(self) -> list[ActiveHook]: + with self._lock: return sorted(self._load_active(), key=lambda x: (-x.descriptor.priority, x.activated_at_s)) + + def resolve(self, context: ExecutionContext) -> list[ActiveHook]: + """Select the highest-priority matching implementation per logical hook group.""" + selected: dict[str, ActiveHook] = {} + for hook in self.active(): + if not hook.descriptor.when.matches(context): continue + group = hook.descriptor.exclusive_group + current = selected.get(group) + if current is None or hook.descriptor.priority > current.descriptor.priority: selected[group] = hook + return sorted(selected.values(), key=lambda x: (-x.descriptor.priority, x.activated_at_s)) + + def _required_status(self, spec: KernelSpec) -> frozenset[str]: + heldout = spec.metadata.get("heldout_command", ()) + return frozenset({"heldout_passed", "accepted"}) if heldout else frozenset({"hardware_passed", "heldout_passed", "accepted"}) + + def activate(self, candidate_id: str, fingerprint: RuntimeFingerprint, reason: str) -> ActiveHook: + if not reason.strip(): raise HookActivationError("activation requires a reason") + with self._lock: + candidate = self.workspace.load_candidate(candidate_id) + spec = self.workspace.load_spec(candidate.spec_id) + descriptor = HookDescriptor.from_spec(spec) + if descriptor.adapter not in self.SUPPORTED_ADAPTERS: raise HookActivationError(f"runtime adapter {descriptor.adapter!r} is not installed") + if candidate.status not in self._required_status(spec): + raise HookActivationError(f"candidate status {candidate.status!r} has not passed the required hardware/held-out oracle") + source_path = Path(candidate.source_path) + if not source_path.is_file(): raise HookActivationError("candidate source is missing") + source_sha = hashlib.sha256(source_path.read_bytes()).hexdigest() + if source_sha != candidate.source_sha256: raise HookActivationError("candidate source hash changed after validation") + compatibility = check_compatibility(spec, fingerprint) + if not compatibility.compatible: raise HookActivationError("incompatible runtime: " + "; ".join(compatibility.mismatches)) + current = self._load_active() + previous = next((x for x in current if x.descriptor.slot_key == descriptor.slot_key), None) + current = [x for x in current if x.descriptor.slot_key != descriptor.slot_key] + active = ActiveHook(uuid.uuid4().hex, candidate.candidate_id, spec.spec_id, candidate.source_path, source_sha, descriptor, + compatibility, time.time(), reason.strip(), previous.activation_id if previous else None) + current.append(active) + self._save_active(current) + self._history("hook_activated", activation=asdict(active), replaced=asdict(previous) if previous else None) + return active + + def deactivate(self, activation_id: str, reason: str) -> ActiveHook: + if not reason.strip(): raise HookActivationError("deactivation requires a reason") + with self._lock: + current = self._load_active() + target = next((x for x in current if x.activation_id == activation_id), None) + if target is None: raise KeyError(f"unknown active hook {activation_id}") + self._save_active([x for x in current if x.activation_id != activation_id]) + self._history("hook_deactivated", activation=asdict(target), reason=reason.strip()) + return target + + def clear(self, reason: str) -> list[ActiveHook]: + with self._lock: + current = self._load_active() + self._save_active([]) + self._history("hooks_cleared", activations=[asdict(x) for x in current], reason=reason.strip()) + return current + + @staticmethod + def _runtime_item(active: ActiveHook) -> dict[str, Any] | None: + source = Path(active.source_path) + if not source.is_file() or hashlib.sha256(source.read_bytes()).hexdigest() != active.source_sha256: return None + return {"activation_id": active.activation_id, "candidate_id": active.candidate_id, "spec_id": active.spec_id, + "source_path": active.source_path, "source_sha256": active.source_sha256, + "descriptor": asdict(active.descriptor)} + + def runtime_metadata(self, context: ExecutionContext | None = None) -> dict[str, Any]: + source = self.resolve(context) if context is not None else self.active() + hooks = [item for active in source if (item := self._runtime_item(active)) is not None] + return {"active_hooks": hooks, "execution_context": asdict(context) if context is not None else None} diff --git a/extra/radeon_forge/synthesis/portable.py b/extra/radeon_forge/synthesis/portable.py new file mode 100644 index 0000000000000..509eb49da4a4c --- /dev/null +++ b/extra/radeon_forge/synthesis/portable.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Mapping + +from .hooks import HookDescriptor +from .recipe import export_recipe +from .workspace import CandidateWorkspace + + +def _toml_value(value: Any) -> str: + if isinstance(value, bool): return "true" if value else "false" + if isinstance(value, (int, float)): return str(value) + if isinstance(value, str): return json.dumps(value, ensure_ascii=False) + if isinstance(value, (list, tuple)): return "[" + ", ".join(_toml_value(x) for x in value) + "]" + raise TypeError(f"unsupported portable hook value: {type(value).__name__}") + + +def _append_table(lines: list[str], name: str, values: Mapping[str, Any]) -> None: + filtered = {key: value for key, value in values.items() if value is not None and value != {}} + if not filtered: return + lines.append(f"\n[{name}]") + for key, value in filtered.items(): + if isinstance(value, Mapping): continue + lines.append(f"{json.dumps(str(key))} = {_toml_value(value)}") + for key, value in filtered.items(): + if isinstance(value, Mapping): _append_table(lines, f"{name}.{key}", value) + + +def export_recipe_with_hook(workspace: CandidateWorkspace, spec_id: str, output: str | Path, + candidate_id: str | None = None) -> Path: + """Export placement, execution-state applicability and empirical search knowledge.""" + path = export_recipe(workspace, spec_id, output, candidate_id) + spec = workspace.load_spec(spec_id) + descriptor = HookDescriptor.from_spec(spec) + predicate = descriptor.when + when = { + "stages": [stage.value for stage in predicate.stages], + "min_context_tokens": predicate.min_context_tokens, + "max_context_tokens": predicate.max_context_tokens, + "min_prompt_tokens": predicate.min_prompt_tokens, + "max_prompt_tokens": predicate.max_prompt_tokens, + "min_generated_token_index": predicate.min_generated_token_index, + "max_generated_token_index": predicate.max_generated_token_index, + "batch_sizes": list(predicate.batch_sizes), + "prefix_cache": predicate.prefix_cache, + "warm": predicate.warm, + "conditions": dict(predicate.conditions), + } + hook = { + "layer": descriptor.layer.value, + "target": descriptor.target, + "mode": descriptor.mode.value, + "adapter": descriptor.adapter, + "priority": descriptor.priority, + "exclusive_group": descriptor.exclusive_group, + "description": descriptor.description, + "selector": dict(descriptor.selector), + "when": when, + } + lines = path.read_text(encoding="utf-8").rstrip().splitlines() + _append_table(lines, "metadata.hook", hook) + + # The search space is portable knowledge, not a mandatory implementation + # abstraction. A receiving agent may reuse, shrink or replace it, but keeping + # it next to the stage predicate makes prior hardware exploration reproducible. + search = spec.metadata.get("search", {}) + if isinstance(search, Mapping) and search: _append_table(lines, "metadata.search", search) + + selected_parameters = {} + if candidate_id is not None: + candidate = workspace.load_candidate(candidate_id) + selected_parameters = candidate.evidence.get("selected_parameters", {}) + if isinstance(selected_parameters, Mapping) and selected_parameters: + _append_table(lines, "metadata.exported_winner", dict(selected_parameters)) + + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + return path diff --git a/extra/radeon_forge/synthesis/recipe.py b/extra/radeon_forge/synthesis/recipe.py new file mode 100644 index 0000000000000..cd69659e554ef --- /dev/null +++ b/extra/radeon_forge/synthesis/recipe.py @@ -0,0 +1,281 @@ +from __future__ import annotations + +import base64, hashlib, json, shutil, tomllib +from dataclasses import asdict, dataclass, field +from pathlib import Path, PurePosixPath +from typing import Any, Mapping, Sequence + +from .workspace import CandidateRecord, CandidateWorkspace, KernelSpec + + +FORMAT_VERSION = 1 +MAX_RECIPE_BYTES = 8_000_000 +MAX_ARTIFACT_BYTES = 4_000_000 +_ALLOWED_ROLES = {"oracle", "benchmark", "heldout", "knowledge", "reference", "seed", "implementation_cache", "support"} + + +def _canonical_hash(payload: Any) -> str: + return hashlib.sha256(json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str).encode()).hexdigest() + + +def _string_list(value: Any, field_name: str) -> tuple[str, ...]: + if value is None: return () + if not isinstance(value, list) or not all(isinstance(x, str) for x in value): raise ValueError(f"{field_name} must be an array of strings") + return tuple(value) + + +def _safe_relative_path(value: str) -> str: + path = PurePosixPath(value) + if not value or path.is_absolute() or ".." in path.parts or "." in path.parts or "\x00" in value: + raise ValueError(f"unsafe recipe artifact path: {value!r}") + return path.as_posix() + + +def _command(value: Any, field_name: str) -> tuple[str, ...]: + command = _string_list(value, field_name) + if any("\x00" in part for part in command): raise ValueError(f"{field_name} contains NUL") + return command + + +@dataclass(frozen=True) +class RecipeArtifact: + path: str + role: str + content: str + executable: bool = False + description: str = "" + + @property + def sha256(self) -> str: return hashlib.sha256(self.content.encode()).hexdigest() + + +@dataclass(frozen=True) +class ForgeRecipe: + """Portable optimization knowledge. + + A recipe deliberately captures intent, invariants, oracles and optional seed + material without defining a kernel DSL. An intelligent executor is expected + to fill implementation gaps and the independent oracles decide what survives. + """ + name: str + description: str + operation: str + target: str + objective: str + agent_brief: str + invariants: tuple[str, ...] + unknowns: tuple[str, ...] = () + compatibility: Mapping[str, Any] = field(default_factory=dict) + acceptance: Mapping[str, Any] = field(default_factory=dict) + shapes: Mapping[str, int | str] = field(default_factory=dict) + dtypes: Mapping[str, str] = field(default_factory=dict) + language: str = "python" + extension: str = ".py" + entrypoint: str = "build_kernel" + mockgpu_command: tuple[str, ...] = () + hardware_command: tuple[str, ...] = () + heldout_command: tuple[str, ...] = () + artifacts: tuple[RecipeArtifact, ...] = () + seed_artifact: str | None = None + seed_hypothesis: str = "Imported implementation cache; revalidate and specialize locally." + metadata: Mapping[str, Any] = field(default_factory=dict) + source_text: str = field(default="", repr=False, compare=False) + + @property + def recipe_id(self) -> str: + payload = asdict(self) + payload.pop("source_text", None) + return f"{self.name}-{_canonical_hash(payload)[:16]}" + + @classmethod + def load(cls, path: str | Path) -> ForgeRecipe: + source_path = Path(path) + raw = source_path.read_bytes() + if len(raw) > MAX_RECIPE_BYTES: raise ValueError(f"recipe exceeds {MAX_RECIPE_BYTES} bytes") + payload = tomllib.loads(raw.decode("utf-8")) + if int(payload.get("format_version", 0)) != FORMAT_VERSION: raise ValueError(f"unsupported recipe format_version; expected {FORMAT_VERSION}") + + required = ("name", "description", "operation", "target", "objective", "agent_brief", "invariants") + missing = [key for key in required if key not in payload] + if missing: raise ValueError(f"recipe missing required fields: {missing}") + + artifacts: list[RecipeArtifact] = [] + seen_paths: set[str] = set() + for item in payload.get("artifact", []): + if not isinstance(item, dict): raise ValueError("each [[artifact]] entry must be a table") + artifact_path = _safe_relative_path(str(item.get("path", ""))) + if artifact_path in seen_paths: raise ValueError(f"duplicate artifact path: {artifact_path}") + seen_paths.add(artifact_path) + role = str(item.get("role", "support")) + if role not in _ALLOWED_ROLES: raise ValueError(f"unsupported artifact role {role!r}") + has_text, has_b64 = "content" in item, "content_base64" in item + if has_text == has_b64: raise ValueError(f"artifact {artifact_path} must define exactly one of content or content_base64") + try: content = str(item["content"]) if has_text else base64.b64decode(str(item["content_base64"]), validate=True).decode("utf-8") + except Exception as exc: raise ValueError(f"artifact {artifact_path} has invalid UTF-8/base64 content") from exc + if len(content.encode()) > MAX_ARTIFACT_BYTES: raise ValueError(f"artifact {artifact_path} exceeds {MAX_ARTIFACT_BYTES} bytes") + expected_sha = item.get("sha256") + actual_sha = hashlib.sha256(content.encode()).hexdigest() + if expected_sha is not None and str(expected_sha) != actual_sha: raise ValueError(f"artifact {artifact_path} sha256 mismatch") + artifacts.append(RecipeArtifact(artifact_path, role, content, bool(item.get("executable", False)), str(item.get("description", "")))) + + seed_artifact = payload.get("seed_artifact") + if seed_artifact is not None: + seed_artifact = _safe_relative_path(str(seed_artifact)) + if seed_artifact not in seen_paths: raise ValueError("seed_artifact must reference an embedded artifact") + + oracle = payload.get("oracle", {}) + if not isinstance(oracle, dict): raise ValueError("[oracle] must be a table") + reserved = {"format_version", "name", "description", "operation", "target", "objective", "agent_brief", "invariants", + "unknowns", "compatibility", "acceptance", "shapes", "dtypes", "language", "extension", "entrypoint", "oracle", + "artifact", "seed_artifact", "seed_hypothesis", "metadata"} + extra = {key: value for key, value in payload.items() if key not in reserved} + metadata = dict(payload.get("metadata", {})) + if extra: metadata["unrecognized_recipe_fields"] = extra + + recipe = cls( + name=str(payload["name"]).strip(), description=str(payload["description"]).strip(), operation=str(payload["operation"]).strip(), + target=str(payload["target"]).strip(), objective=str(payload["objective"]).strip(), agent_brief=str(payload["agent_brief"]).strip(), + invariants=_string_list(payload["invariants"], "invariants"), unknowns=_string_list(payload.get("unknowns", []), "unknowns"), + compatibility=dict(payload.get("compatibility", {})), acceptance=dict(payload.get("acceptance", {})), + shapes=dict(payload.get("shapes", {})), dtypes={str(k): str(v) for k, v in dict(payload.get("dtypes", {})).items()}, + language=str(payload.get("language", "python")), extension=str(payload.get("extension", ".py")), + entrypoint=str(payload.get("entrypoint", "build_kernel")), + mockgpu_command=_command(oracle.get("mockgpu_command", []), "oracle.mockgpu_command"), + hardware_command=_command(oracle.get("hardware_command", []), "oracle.hardware_command"), + heldout_command=_command(oracle.get("heldout_command", []), "oracle.heldout_command"), artifacts=tuple(artifacts), + seed_artifact=seed_artifact, seed_hypothesis=str(payload.get("seed_hypothesis", "Imported implementation cache; revalidate and specialize locally.")), + metadata=metadata, source_text=raw.decode("utf-8"), + ) + if not recipe.name or not recipe.agent_brief or not recipe.invariants: raise ValueError("name, agent_brief and at least one invariant are required") + if not recipe.extension.startswith(".") or "/" in recipe.extension: raise ValueError("extension must be a simple suffix such as .py or .cpp") + return recipe + + +@dataclass(frozen=True) +class InstalledRecipe: + recipe_id: str + bundle_root: str + spec_id: str + seed_candidate_id: str | None + artifact_hashes: Mapping[str, str] + + +class RecipeLibrary: + def __init__(self, workspace: CandidateWorkspace): + self.workspace = workspace + self.root = workspace.root / "recipes" + self.root.mkdir(parents=True, exist_ok=True) + + def install(self, recipe: ForgeRecipe) -> InstalledRecipe: + bundle = self.root / recipe.recipe_id + if bundle.exists(): shutil.rmtree(bundle) + bundle.mkdir(parents=True) + if recipe.source_text: (bundle / "recipe.forge.toml").write_text(recipe.source_text, encoding="utf-8") + artifact_hashes: dict[str, str] = {} + for artifact in recipe.artifacts: + path = (bundle / artifact.path).resolve() + if bundle.resolve() not in path.parents: raise ValueError(f"artifact escaped bundle: {artifact.path}") + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(artifact.content, encoding="utf-8") + if artifact.executable: path.chmod(path.stat().st_mode | 0o100) + artifact_hashes[artifact.path] = artifact.sha256 + + metadata = { + **dict(recipe.metadata), "recipe_id": recipe.recipe_id, "recipe_bundle": str(bundle), "recipe_agent_brief": recipe.agent_brief, + "recipe_unknowns": list(recipe.unknowns), "recipe_compatibility": dict(recipe.compatibility), "recipe_acceptance": dict(recipe.acceptance), + "recipe_artifacts": [{"path": x.path, "role": x.role, "sha256": x.sha256, "description": x.description} for x in recipe.artifacts], + "heldout_command": list(recipe.heldout_command), + } + spec = KernelSpec(recipe.name, recipe.operation, recipe.target, recipe.language, recipe.extension, recipe.entrypoint, + recipe.shapes, recipe.dtypes, recipe.invariants, recipe.objective, recipe.mockgpu_command, + recipe.hardware_command, metadata) + self.workspace.save_spec(spec) + + seed_candidate: CandidateRecord | None = None + if recipe.seed_artifact is not None: + source = (bundle / recipe.seed_artifact).read_text(encoding="utf-8") + seed_candidate = self.workspace.create_candidate(spec.spec_id, source, recipe.seed_hypothesis) + seed_candidate = self.workspace.update(seed_candidate.candidate_id, "imported_unverified", { + "recipe_id": recipe.recipe_id, "seed_artifact": recipe.seed_artifact, "portable_cache": True, + }) + + manifest = {"recipe_id": recipe.recipe_id, "spec_id": spec.spec_id, "bundle_root": str(bundle), + "seed_candidate_id": seed_candidate.candidate_id if seed_candidate else None, "artifact_hashes": artifact_hashes} + (bundle / "installed.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8") + return InstalledRecipe(**manifest) + + def install_file(self, path: str | Path) -> InstalledRecipe: return self.install(ForgeRecipe.load(path)) + + def installed(self) -> list[InstalledRecipe]: + ret: list[InstalledRecipe] = [] + for path in sorted(self.root.glob("*/installed.json")): + try: ret.append(InstalledRecipe(**json.loads(path.read_text(encoding="utf-8")))) + except Exception: continue + return ret + + def inspect(self, recipe_id: str) -> dict[str, Any]: + path = self.root / recipe_id / "installed.json" + if not path.is_file(): raise KeyError(f"unknown recipe {recipe_id}") + installed = InstalledRecipe(**json.loads(path.read_text(encoding="utf-8"))) + spec = self.workspace.load_spec(installed.spec_id) + return {"installed": asdict(installed), "spec": asdict(spec) | {"spec_id": spec.spec_id}, + "recipe_source": str(self.root / recipe_id / "recipe.forge.toml")} + + +def _toml_string(value: str) -> str: return json.dumps(value, ensure_ascii=False) +def _toml_array(values: Sequence[str]) -> str: return "[" + ", ".join(_toml_string(x) for x in values) + "]" + + +def export_recipe(workspace: CandidateWorkspace, spec_id: str, output: str | Path, candidate_id: str | None = None) -> Path: + """Export a spec and optional implementation cache as one portable TOML file. + + Exact generated code is base64-encoded because it is a disposable cache. The + intent, invariants and agent brief stay human-readable and reviewable. + """ + spec = workspace.load_spec(spec_id) + metadata = dict(spec.metadata) + agent_brief = str(metadata.get("recipe_agent_brief") or + "Use the supplied intent, invariants and oracle as the contract. Inspect the local hardware and workload, then freely regenerate or restructure the implementation. Do not trust the cached implementation without rerunning every oracle.") + lines = [f"format_version = {FORMAT_VERSION}", f"name = {_toml_string(spec.name)}", + f"description = {_toml_string(str(metadata.get('description', spec.operation)))}", f"operation = {_toml_string(spec.operation)}", + f"target = {_toml_string(spec.target)}", f"objective = {_toml_string(spec.objective)}", f"agent_brief = {_toml_string(agent_brief)}", + f"invariants = {_toml_array(spec.invariants)}", f"language = {_toml_string(spec.language)}", + f"extension = {_toml_string(spec.extension)}", f"entrypoint = {_toml_string(spec.entrypoint)}"] + unknowns = tuple(str(x) for x in metadata.get("recipe_unknowns", ())) + if unknowns: lines.append(f"unknowns = {_toml_array(unknowns)}") + if candidate_id is not None: + candidate = workspace.load_candidate(candidate_id) + if candidate.spec_id != spec_id: raise ValueError("candidate does not belong to spec") + source = Path(candidate.source_path).read_text(encoding="utf-8") + seed_name = f"implementation-cache{spec.extension}" + lines += [f"seed_artifact = {_toml_string(seed_name)}", f"seed_hypothesis = {_toml_string(candidate.hypothesis)}"] + else: candidate, source, seed_name = None, "", "" + + if spec.shapes: + lines.append("\n[shapes]") + lines += [f"{json.dumps(str(k))} = {json.dumps(v)}" for k, v in spec.shapes.items()] + if spec.dtypes: + lines.append("\n[dtypes]") + lines += [f"{json.dumps(str(k))} = {_toml_string(str(v))}" for k, v in spec.dtypes.items()] + compatibility = dict(metadata.get("recipe_compatibility", {})) + if compatibility: + lines.append("\n[compatibility]") + lines += [f"{json.dumps(str(k))} = {json.dumps(v)}" for k, v in compatibility.items()] + acceptance = dict(metadata.get("recipe_acceptance", {})) + if acceptance: + lines.append("\n[acceptance]") + lines += [f"{json.dumps(str(k))} = {json.dumps(v)}" for k, v in acceptance.items()] + lines += ["\n[oracle]", f"mockgpu_command = {_toml_array(spec.mockgpu_command)}", f"hardware_command = {_toml_array(spec.hardware_command)}"] + heldout = tuple(str(x) for x in metadata.get("heldout_command", ())) + if heldout: lines.append(f"heldout_command = {_toml_array(heldout)}") + + if candidate is not None: + encoded = base64.b64encode(source.encode()).decode() + lines += ["\n[[artifact]]", f"path = {_toml_string(seed_name)}", "role = \"implementation_cache\"", + f"description = {_toml_string('Disposable implementation cache exported from ' + candidate.candidate_id)}", + f"sha256 = {_toml_string(hashlib.sha256(source.encode()).hexdigest())}", f"content_base64 = {_toml_string(encoded)}"] + + out = Path(output) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text("\n".join(lines) + "\n", encoding="utf-8") + return out diff --git a/extra/radeon_forge/synthesis/scaffold.py b/extra/radeon_forge/synthesis/scaffold.py new file mode 100644 index 0000000000000..1dc7da04e7011 --- /dev/null +++ b/extra/radeon_forge/synthesis/scaffold.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import json +from typing import Any + +from .hooks import HookDescriptor, HookLayer +from .workspace import KernelSpec + + +def _doc(value: Any) -> str: return json.dumps(value, indent=2, sort_keys=True, default=str) + + +def render_candidate_scaffold(spec: KernelSpec) -> str: + """Render the minimum executable contract, not a performance abstraction. + + The scaffold intentionally contains no tile DSL or scheduling framework. It + gives an agent the exact runtime boundary, invariants, stage predicate and + parameter handoff, while leaving the implementation structure disposable. + """ + hook = HookDescriptor.from_spec(spec) + header = f'''"""Disposable Radeon Forge candidate. + +Spec: {spec.name} +Operation: {spec.operation} +Target: {spec.target} +Objective: {spec.objective} +Placement: {hook.layer.value}:{hook.target} ({hook.mode.value}) +Execution states: {hook.when.signature} + +Invariants: +{chr(10).join(f"- {item}" for item in spec.invariants)} + +This file is not trusted merely because it imports or compiles. It must pass +MockGPU/reference, real hardware, and any held-out phase-transition oracle. +""" + +from __future__ import annotations +from typing import Any, Mapping + +RADEON_FORGE_PARAMETERS: dict[str, Any] = {{}} + + +def configure(parameters: Mapping[str, Any]) -> None: + """Receive the locally autotuned schedule selected for this machine/state.""" + global RADEON_FORGE_PARAMETERS + RADEON_FORGE_PARAMETERS = dict(parameters) + +''' + if hook.layer is HookLayer.TRANSFORMER_BLOCK: + body = '''def build_replacement(original, layer_index: int, model, context: Mapping[str, Any]): + """Return a callable with signature (x, start_pos, freqs_cis, mask). + + `context` describes the activation/validation state. Runtime execution-state + selection is enforced outside this module by the signed hook manifest. + Replace this identity implementation with direct tinygrad/custom-kernel code. + """ + assert context.get("stage") in {"transition_oracle", "decode_benchmark", "first_token", "decode"} + + def replacement(x, start_pos, freqs_cis, mask): + # TODO(agent): implement the model- and gfx1100-specific block/subgraph. + # Keep `original` as the correctness fallback while iterating. + return original(x, start_pos, freqs_cis, mask) + + return replacement +''' + else: + body = '''def build_kernel(*args, **kwargs): + """Build the direct target-specific implementation for this hook contract.""" + raise NotImplementedError("agent must generate a direct implementation") +''' + footer = f'''\n\n# Machine-readable context copied from the durable spec for code-review tools.\nFORGE_SPEC_CONTEXT = {_doc({ + "shapes":dict(spec.shapes), "dtypes":dict(spec.dtypes), "invariants":list(spec.invariants), + "hook":{"layer":hook.layer.value, "target":hook.target, "mode":hook.mode.value, + "adapter":hook.adapter, "selector":dict(hook.selector), "when":hook.when.signature}, + "search":spec.metadata.get("search", {}), + })}\n''' + return header + body + footer diff --git a/extra/radeon_forge/synthesis/tools.py b/extra/radeon_forge/synthesis/tools.py new file mode 100644 index 0000000000000..4b4ee8359c626 --- /dev/null +++ b/extra/radeon_forge/synthesis/tools.py @@ -0,0 +1,155 @@ +from __future__ import annotations + +import os, subprocess, time +from dataclasses import asdict +from pathlib import Path +from typing import Any, Callable, Mapping + +from ..oracles.mockgpu import MockGPUOracle +from ..permissions import Action +from ..runtime.tools import ToolRegistry, ToolSpec +from .autotune import SearchPlan, run_autotune +from .hooks import HookRegistry, RuntimeFingerprint, check_compatibility +from .portable import export_recipe_with_hook +from .recipe import RecipeLibrary +from .workspace import CandidateWorkspace + + +class OptimizationTools: + def __init__(self, workspace: CandidateWorkspace, project_root: str | Path, hooks: HookRegistry | None = None, + fingerprint_provider: Callable[[], RuntimeFingerprint] | None = None): + self.workspace = workspace + self.project_root = Path(project_root).resolve() + self.mockgpu = MockGPUOracle(self.project_root) + self.recipes = RecipeLibrary(workspace) + self.hooks, self.fingerprint_provider = hooks, fingerprint_provider + + def _project_path(self, value: str, *, must_exist: bool = False) -> Path: + path = (self.project_root / value).resolve() if not Path(value).is_absolute() else Path(value).resolve() + if path != self.project_root and self.project_root not in path.parents: raise ValueError("path must stay inside the private workspace") + if must_exist and not path.is_file(): raise FileNotFoundError(path) + return path + + def _fingerprint(self) -> RuntimeFingerprint: + return self.fingerprint_provider() if self.fingerprint_provider is not None else RuntimeFingerprint() + + def list_specs(self, args: Mapping[str, Any]) -> Any: + return {"specs": [asdict(x) | {"spec_id": x.spec_id} for x in self.workspace.specs()]} + + def list_candidates(self, args: Mapping[str, Any]) -> Any: + return {"candidates": [asdict(x) for x in self.workspace.candidates(args.get("spec_id"))]} + + def stage_candidate(self, args: Mapping[str, Any]) -> Any: + record = self.workspace.create_candidate(str(args["spec_id"]), str(args["source"]), str(args["hypothesis"]), args.get("parent_id")) + return asdict(record) + + def validate_mockgpu(self, args: Mapping[str, Any]) -> Any: + candidate = self.workspace.load_candidate(str(args["candidate_id"])) + spec = self.workspace.load_spec(candidate.spec_id) + if not spec.mockgpu_command: raise ValueError("spec has no MockGPU oracle command") + command = self.workspace.render_command(spec.mockgpu_command, spec, candidate, self.project_root) + result = self.mockgpu.run(command, env={"RADEON_FORGE_CANDIDATE": candidate.source_path, + "RADEON_FORGE_RECIPE_BUNDLE": str(spec.metadata.get("recipe_bundle", ""))}) + updated = self.workspace.update(candidate.candidate_id, "mockgpu_passed" if result.passed else "mockgpu_failed", {"mockgpu": result.to_dict()}) + return asdict(updated) + + def _run_hardware_command(self, candidate_id: str, command_value, timeout_seconds: int, evidence_key: str, + pass_status: str, fail_status: str) -> Any: + candidate = self.workspace.load_candidate(candidate_id) + spec = self.workspace.load_spec(candidate.spec_id) + command = self.workspace.render_command(command_value, spec, candidate, self.project_root) + env = {**os.environ, "DEV": "AMD", "RADEON_FORGE_CANDIDATE": candidate.source_path, + "RADEON_FORGE_RECIPE_BUNDLE": str(spec.metadata.get("recipe_bundle", "")), "PYTHONUNBUFFERED": "1"} + started = time.perf_counter_ns() + proc = subprocess.run(command, cwd=str(spec.metadata.get("recipe_bundle", self.project_root)), env=env, text=True, capture_output=True, + timeout=timeout_seconds) + evidence = {"command": command, "returncode": proc.returncode, "elapsed_ms": (time.perf_counter_ns()-started)/1e6, + "stdout": proc.stdout[-50000:], "stderr": proc.stderr[-50000:], "target": spec.target} + updated = self.workspace.update(candidate.candidate_id, pass_status if proc.returncode == 0 else fail_status, {evidence_key: evidence}) + return asdict(updated) + + def benchmark_hardware(self, args: Mapping[str, Any]) -> Any: + candidate = self.workspace.load_candidate(str(args["candidate_id"])) + if candidate.status != "mockgpu_passed" and not bool(args.get("allow_without_mockgpu", False)): + raise ValueError("candidate must pass MockGPU before W7900 benchmarking") + spec = self.workspace.load_spec(candidate.spec_id) + if not spec.hardware_command: raise ValueError("spec has no hardware benchmark command") + return self._run_hardware_command(candidate.candidate_id, spec.hardware_command, int(args.get("timeout_seconds", 900)), + "hardware", "hardware_passed", "hardware_failed") + + def autotune_hardware(self, args: Mapping[str, Any]) -> Any: + candidate = self.workspace.load_candidate(str(args["candidate_id"])) + spec = self.workspace.load_spec(candidate.spec_id) + stored = spec.metadata.get("search", {}) + if not isinstance(stored, Mapping): stored = {} + override = {key: args[key] for key in ("axes", "budgets", "reduction", "max_candidates", "timeout_seconds") if key in args} + plan = SearchPlan.from_mapping({**dict(stored), **override}) + updated, summary = run_autotune(self.workspace, candidate.candidate_id, self.project_root, plan) + return {"candidate": asdict(updated), "winner": summary.winner.to_dict() if summary.winner else None, + "rounds": summary.rounds, "evaluated_trials": len(summary.evaluated)} + + def validate_heldout(self, args: Mapping[str, Any]) -> Any: + candidate = self.workspace.load_candidate(str(args["candidate_id"])) + if candidate.status != "hardware_passed": raise ValueError("candidate must pass the real hardware benchmark or autotuner first") + spec = self.workspace.load_spec(candidate.spec_id) + command = tuple(str(x) for x in spec.metadata.get("heldout_command", ())) + if not command: raise ValueError("spec has no held-out validation command") + return self._run_hardware_command(candidate.candidate_id, command, int(args.get("timeout_seconds", 900)), + "heldout", "heldout_passed", "heldout_failed") + + def list_recipes(self, args: Mapping[str, Any]) -> Any: + return {"recipes": [asdict(x) for x in self.recipes.installed()]} + + def inspect_recipe(self, args: Mapping[str, Any]) -> Any: + return self.recipes.inspect(str(args["recipe_id"])) + + def import_recipe(self, args: Mapping[str, Any]) -> Any: + path = self._project_path(str(args["path"]), must_exist=True) + return asdict(self.recipes.install_file(path)) + + def export_recipe_file(self, args: Mapping[str, Any]) -> Any: + output = self._project_path(str(args["output"])) + path = export_recipe_with_hook(self.workspace, str(args["spec_id"]), output, + str(args["candidate_id"]) if args.get("candidate_id") else None) + return {"path": str(path), "sha256": __import__("hashlib").sha256(path.read_bytes()).hexdigest(), "portable": True, + "includes_execution_stage_hook": True} + + def list_active_hooks(self, args: Mapping[str, Any]) -> Any: + if self.hooks is None: return {"hooks": [], "available": False} + return {"hooks": [asdict(x) for x in self.hooks.active()], "available": True, + "runtime_fingerprint": asdict(self._fingerprint())} + + def activate_hook(self, args: Mapping[str, Any]) -> Any: + if self.hooks is None: raise RuntimeError("hook registry is unavailable") + candidate = self.workspace.load_candidate(str(args["candidate_id"])) + spec = self.workspace.load_spec(candidate.spec_id) + fingerprint = self._fingerprint() + allow_unknown = bool(args.get("allow_unknown_runtime", False)) + report = check_compatibility(spec, fingerprint) + if report.mismatches: raise ValueError("runtime is incompatible: " + "; ".join(report.mismatches)) + if report.unknown and not allow_unknown: + raise ValueError("runtime compatibility is incomplete: missing " + ", ".join(report.unknown) + + ". Regenerate locally or explicitly approve an unknown-runtime deployment.") + active = self.hooks.activate(candidate.candidate_id, fingerprint, + str(args.get("reason", "Approved stage-specific local deployment"))) + return asdict(active) + + def deactivate_hook(self, args: Mapping[str, Any]) -> Any: + if self.hooks is None: raise RuntimeError("hook registry is unavailable") + return asdict(self.hooks.deactivate(str(args["activation_id"]), str(args.get("reason", "User-requested rollback")))) + + def install(self, registry: ToolRegistry) -> None: + registry.register(ToolSpec("list_kernel_specs", "List typed megakernel and fused-kernel contracts", {"type":"object","properties":{}}), self.list_specs) + registry.register(ToolSpec("list_kernel_candidates", "List staged generated implementations and their oracle status", {"type":"object","properties":{"spec_id":{"type":"string"}}}), self.list_candidates) + registry.register(ToolSpec("stage_kernel_candidate", "Write one disposable implementation for a typed kernel specification", {"type":"object","required":["spec_id","source","hypothesis"],"properties":{"spec_id":{"type":"string"},"source":{"type":"string"},"hypothesis":{"type":"string"},"parent_id":{"type":"string"}}}, Action.WRITE_GENERATED_SOURCE), self.stage_candidate) + registry.register(ToolSpec("validate_candidate_mockgpu", "Execute a generated AMD candidate through tinygrad's RDNA3 MockGPU semantic oracle", {"type":"object","required":["candidate_id"],"properties":{"candidate_id":{"type":"string"}}}, Action.COMPILE), self.validate_mockgpu) + registry.register(ToolSpec("benchmark_candidate_w7900", "Benchmark a MockGPU-passing candidate on the real local W7900", {"type":"object","required":["candidate_id"],"properties":{"candidate_id":{"type":"string"},"timeout_seconds":{"type":"integer"},"allow_without_mockgpu":{"type":"boolean"}}}, Action.BENCHMARK), self.benchmark_hardware) + registry.register(ToolSpec("autotune_candidate_w7900", "Empirically search a bounded parameter space for one MockGPU-passing implementation in its declared execution stages", {"type":"object","required":["candidate_id","axes"],"properties":{"candidate_id":{"type":"string"},"axes":{"type":"object"},"budgets":{"type":"array"},"reduction":{"type":"integer"},"max_candidates":{"type":"integer"},"timeout_seconds":{"type":"integer"}}}, Action.BENCHMARK), self.autotune_hardware) + registry.register(ToolSpec("validate_candidate_heldout", "Run a hardware-passing candidate on its held-out correctness and task-quality suite", {"type":"object","required":["candidate_id"],"properties":{"candidate_id":{"type":"string"},"timeout_seconds":{"type":"integer"}}}, Action.BENCHMARK), self.validate_heldout) + registry.register(ToolSpec("list_forge_recipes", "List portable installed optimization recipes", {"type":"object","properties":{}}), self.list_recipes) + registry.register(ToolSpec("inspect_forge_recipe", "Read the intent, invariants, execution-stage hook, oracle and artifact inventory of an installed optimization recipe", {"type":"object","required":["recipe_id"],"properties":{"recipe_id":{"type":"string"}}}), self.inspect_recipe) + registry.register(ToolSpec("import_forge_recipe", "Install one portable .forge.toml optimization recipe from the private workspace", {"type":"object","required":["path"],"properties":{"path":{"type":"string"}}}, Action.WRITE_GENERATED_SOURCE), self.import_recipe) + registry.register(ToolSpec("export_forge_recipe", "Export a contract, execution-stage hook and optional implementation cache as one shareable .forge.toml file", {"type":"object","required":["spec_id","output"],"properties":{"spec_id":{"type":"string"},"candidate_id":{"type":"string"},"output":{"type":"string"}}}, Action.WRITE_GENERATED_SOURCE), self.export_recipe_file) + registry.register(ToolSpec("list_active_optimization_hooks", "List deployed optimizations and the execution states in which each can run", {"type":"object","properties":{}}), self.list_active_hooks) + registry.register(ToolSpec("activate_optimization_hook", "Deploy a validated optimization into its declared execution stages; activation is compatibility-gated and rollback-safe", {"type":"object","required":["candidate_id"],"properties":{"candidate_id":{"type":"string"},"reason":{"type":"string"},"allow_unknown_runtime":{"type":"boolean"}}}, Action.DEPLOY), self.activate_hook) + registry.register(ToolSpec("deactivate_optimization_hook", "Rollback one active optimization hook", {"type":"object","required":["activation_id"],"properties":{"activation_id":{"type":"string"},"reason":{"type":"string"}}}, Action.DEPLOY), self.deactivate_hook) diff --git a/extra/radeon_forge/synthesis/workspace.py b/extra/radeon_forge/synthesis/workspace.py new file mode 100644 index 0000000000000..04b36a134f560 --- /dev/null +++ b/extra/radeon_forge/synthesis/workspace.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +import hashlib, json, time +from dataclasses import asdict, dataclass, field, replace +from pathlib import Path +from typing import Any, Mapping, Sequence + + +def _hash(payload: Any) -> str: + return hashlib.sha256(json.dumps(payload, sort_keys=True, default=str).encode()).hexdigest() + + +@dataclass(frozen=True) +class KernelSpec: + name: str + operation: str + target: str = "gfx1100" + language: str = "python" + extension: str = ".py" + entrypoint: str = "build_kernel" + shapes: Mapping[str, int | str] = field(default_factory=dict) + dtypes: Mapping[str, str] = field(default_factory=dict) + invariants: tuple[str, ...] = () + objective: str = "minimize_p95_latency" + mockgpu_command: tuple[str, ...] = () + hardware_command: tuple[str, ...] = () + metadata: Mapping[str, Any] = field(default_factory=dict) + + @property + def spec_id(self) -> str: return f"{self.name}-{_hash(asdict(self))[:12]}" + + +@dataclass(frozen=True) +class CandidateRecord: + candidate_id: str + spec_id: str + source_path: str + source_sha256: str + parent_id: str | None + hypothesis: str + status: str + created_at_s: float + evidence: Mapping[str, Any] = field(default_factory=dict) + + +class CandidateWorkspace: + """Content-addressed store where generated implementations are disposable artifacts.""" + def __init__(self, root: str | Path): + self.root = Path(root).resolve() + self.spec_root, self.candidate_root = self.root / "specs", self.root / "candidates" + self.spec_root.mkdir(parents=True, exist_ok=True) + self.candidate_root.mkdir(parents=True, exist_ok=True) + + def save_spec(self, spec: KernelSpec) -> KernelSpec: + path = self.spec_root / f"{spec.spec_id}.json" + path.write_text(json.dumps(asdict(spec), indent=2, default=str), encoding="utf-8") + return spec + + def load_spec(self, spec_id: str) -> KernelSpec: + payload = json.loads((self.spec_root / f"{spec_id}.json").read_text(encoding="utf-8")) + for key in ("invariants", "mockgpu_command", "hardware_command"): payload[key] = tuple(payload.get(key, ())) + return KernelSpec(**payload) + + def specs(self) -> list[KernelSpec]: + ret = [] + for path in sorted(self.spec_root.glob("*.json")): + try: ret.append(self.load_spec(path.stem)) + except Exception: continue + return ret + + def create_candidate(self, spec_id: str, source: str, hypothesis: str, parent_id: str | None = None) -> CandidateRecord: + spec = self.load_spec(spec_id) + if not source.strip(): raise ValueError("candidate source must not be empty") + if len(source.encode()) > 2_000_000: raise ValueError("candidate source exceeds 2 MB") + source_sha = hashlib.sha256(source.encode()).hexdigest() + candidate_id = f"{spec.name}-{source_sha[:12]}" + directory = self.candidate_root / candidate_id + directory.mkdir(parents=True, exist_ok=True) + source_path = directory / f"candidate{spec.extension}" + source_path.write_text(source, encoding="utf-8") + record = CandidateRecord(candidate_id, spec_id, str(source_path), source_sha, parent_id, hypothesis.strip(), "staged", time.time()) + (directory / "manifest.json").write_text(json.dumps(asdict(record), indent=2, default=str), encoding="utf-8") + return record + + def load_candidate(self, candidate_id: str) -> CandidateRecord: + payload = json.loads((self.candidate_root / candidate_id / "manifest.json").read_text(encoding="utf-8")) + return CandidateRecord(**payload) + + def candidates(self, spec_id: str | None = None) -> list[CandidateRecord]: + ret = [] + for path in sorted(self.candidate_root.glob("*/manifest.json")): + try: + record = CandidateRecord(**json.loads(path.read_text(encoding="utf-8"))) + if spec_id is None or record.spec_id == spec_id: ret.append(record) + except Exception: continue + return sorted(ret, key=lambda x: x.created_at_s, reverse=True) + + def update(self, candidate_id: str, status: str, evidence: Mapping[str, Any]) -> CandidateRecord: + record = self.load_candidate(candidate_id) + merged = {**dict(record.evidence), **dict(evidence)} + updated = replace(record, status=status, evidence=merged) + path = self.candidate_root / candidate_id / "manifest.json" + path.write_text(json.dumps(asdict(updated), indent=2, default=str), encoding="utf-8") + return updated + + @staticmethod + def render_command(command: Sequence[str], spec: KernelSpec, candidate: CandidateRecord, root: Path) -> tuple[str, ...]: + bundle = str(spec.metadata.get("recipe_bundle", root)) + replacements = {"{candidate}": candidate.source_path, "{candidate_id}": candidate.candidate_id, + "{spec_id}": spec.spec_id, "{root}": str(root), "{bundle}": bundle} + rendered: list[str] = [] + for original in command: + part = str(original) + for key, value in replacements.items(): part = part.replace(key, value) + rendered.append(part) + return tuple(rendered) diff --git a/extra/radeon_forge/tuner.py b/extra/radeon_forge/tuner.py new file mode 100644 index 0000000000000..be5414e7aa86b --- /dev/null +++ b/extra/radeon_forge/tuner.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from math import ceil + +from .contracts import Candidate, TrialResult, WorkloadContract +from .ledger import ExperimentLedger + + +Evaluator = Callable[[Candidate, int], TrialResult] + + +@dataclass(frozen=True) +class TuningSummary: + winner: TrialResult | None + evaluated: tuple[TrialResult, ...] + rounds: int + + +def _rank(results: Sequence[TrialResult], contract: WorkloadContract) -> list[TrialResult]: + return sorted(results, key=lambda result: (not result.is_feasible(contract), result.metric(contract.objective), result.candidate.candidate_id)) + + +def successive_halving(candidates: Sequence[Candidate], evaluator: Evaluator, contract: WorkloadContract, budgets: Sequence[int] = (5, 20, 100), + reduction: int = 4, ledger: ExperimentLedger | None = None) -> TuningSummary: + """Evaluate candidates with increasing benchmark budgets. + + Correctness and resource constraints are hard gates. A faster candidate can + never outrank a candidate that passes the workload contract. + """ + if not candidates: return TuningSummary(None, (), 0) + if not budgets or any(budget <= 0 for budget in budgets): raise ValueError("budgets must contain positive iteration counts") + if reduction < 2: raise ValueError("reduction must be at least 2") + + active = list(candidates) + all_results: list[TrialResult] = [] + completed_rounds = 0 + for round_index, budget in enumerate(budgets): + completed_rounds += 1 + round_results: list[TrialResult] = [] + for candidate in active: + if ledger is not None: ledger.append("trial_started", {"round": round_index, "budget": budget, "candidate": candidate}) + try: + result = evaluator(candidate, budget) + except Exception as exc: # the optimizer records failures rather than silently losing candidates + from .contracts import CorrectnessReport + result = TrialResult(candidate=candidate, correctness=CorrectnessReport(False, reason=f"evaluator exception: {exc!r}"), + compile_ok=False, stable=False, rejected_reason="evaluator_exception") + round_results.append(result) + all_results.append(result) + if ledger is not None: ledger.append("trial_finished", {"round": round_index, "budget": budget, "result": result}) + + ranked = _rank(round_results, contract) + feasible = [result for result in ranked if result.is_feasible(contract)] + if not feasible: + if ledger is not None: ledger.append("round_failed", {"round": round_index, "reason": "no feasible candidates"}) + return TuningSummary(None, tuple(all_results), completed_rounds) + + if round_index == len(budgets) - 1: + winner = feasible[0] + if ledger is not None: ledger.append("winner_selected", winner) + return TuningSummary(winner, tuple(all_results), completed_rounds) + + keep = max(1, ceil(len(feasible) / reduction)) + active = [result.candidate for result in feasible[:keep]] + if ledger is not None: ledger.append("round_survivors", {"round": round_index, "candidate_ids": [candidate.candidate_id for candidate in active]}) + + raise AssertionError("unreachable") diff --git a/extra/radeon_forge/ui/openai_api.py b/extra/radeon_forge/ui/openai_api.py new file mode 100644 index 0000000000000..0ff666b56d17e --- /dev/null +++ b/extra/radeon_forge/ui/openai_api.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +import json, statistics, time, uuid +from collections import defaultdict +from typing import Any, Iterable, Mapping, Sequence + +from ..runtime import ForgeEngine, GenerationEvent + + +_TOOL_PREFIX = "" + + +class OpenAIRequestError(ValueError): pass + + +def _request(payload: Mapping[str, Any]) -> tuple[list[Mapping[str, Any]], list[Mapping[str, Any]], int, float, str, tuple[str, ...]]: + messages = payload.get("messages") + if not isinstance(messages, list) or not messages or not all(isinstance(item, Mapping) for item in messages): + raise OpenAIRequestError("messages must be a non-empty array of objects") + tools = payload.get("tools", []) + if not isinstance(tools, list) or not all(isinstance(item, Mapping) for item in tools): raise OpenAIRequestError("tools must be an array") + max_tokens = int(payload.get("max_tokens", payload.get("max_completion_tokens", 256))) + if max_tokens <= 0: raise OpenAIRequestError("max_tokens must be positive") + temperature = float(payload.get("temperature", 0.0)) + session_id = str(payload.get("session_id") or payload.get("user") or uuid.uuid4().hex) + stop = payload.get("stop", ()) + if isinstance(stop, str): stop = (stop,) + elif isinstance(stop, list) and all(isinstance(item, str) for item in stop): stop = tuple(stop) + elif stop in (None, ()): stop = () + else: raise OpenAIRequestError("stop must be a string or array of strings") + return list(messages), list(tools), max_tokens, temperature, session_id, tuple(stop) + + +def _tool_call(value: Mapping[str, Any]) -> dict[str, Any]: + return {"id": str(value.get("id") or uuid.uuid4().hex), "type": "function", + "function": {"name": str(value["name"]), "arguments": json.dumps(value.get("arguments", {}), separators=(",", ":"))}} + + +def _finish_reason(value: str | None) -> str | None: return "tool_calls" if value == "tool_call" else value + + +def _summarize(events: Sequence[GenerationEvent], session_id: str) -> dict[str, Any]: + prefill = next((dict(event.metrics) for event in events if event.kind == "prefill"), {}) + tokens = [dict(event.metrics) for event in events if event.kind == "token"] + by_stage: dict[str, list[float]] = defaultdict(list) + kernel_calls: dict[str, int] = defaultdict(int) + hook_events = [] + for event in events: + if event.kind == "token": + wall = event.metrics.get("wall_ms") + if wall is not None: by_stage[str(event.metrics.get("stage", "decode"))].append(float(wall)) + elif event.kind == "kernel": kernel_calls[str(event.metrics.get("stage", "unknown"))] += 1 + elif event.kind == "hook": hook_events.append(dict(event.metrics)) + return {"session_id": session_id, "prefill": prefill, + "token_latency_ms": {stage: {"count": len(values), "p50": statistics.median(values), "max": max(values)} + for stage, values in by_stage.items() if values}, + "kernel_calls_by_stage": dict(kernel_calls), "hook_transitions": hook_events, + "generated_tokens": len(tokens)} + + +def collect_chat_completion(engine: ForgeEngine, payload: Mapping[str, Any]) -> dict[str, Any]: + messages, tools, max_tokens, temperature, session_id, stop = _request(payload) + events = list(engine.stream_inference(messages, tools, max_tokens, temperature, session_id, stop)) + text = "".join(event.text for event in events if event.kind == "token") + tool = next((event.tool_call for event in events if event.kind == "tool_call" and event.tool_call is not None), None) + done = next((event for event in reversed(events) if event.kind == "done"), None) + prefill = next((event for event in events if event.kind == "prefill"), None) + completion_tokens = int(done.metrics.get("generated_tokens", 0)) if done is not None else sum(event.kind == "token" for event in events) + prompt_tokens = int(prefill.metrics.get("prompt_tokens", 0)) if prefill is not None else 0 + message: dict[str, Any] = {"role": "assistant", "content": None if tool is not None else text} + if tool is not None: message["tool_calls"] = [_tool_call(tool)] + return {"id": f"chatcmpl-{uuid.uuid4().hex}", "object": "chat.completion", "created": int(time.time()), + "model": str(payload.get("model") or engine.backend.name), + "choices": [{"index": 0, "message": message, "finish_reason": _finish_reason(done.finish_reason if done else None)}], + "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens}, + "system_fingerprint": engine.runtime_fingerprint().architecture or None, + "forge": _summarize(events, session_id)} + + +def _content_chunk(completion_id: str, created: int, model: str, text: str, include_role: bool) -> str: + delta = {"content": text} + if include_role: delta = {"role": "assistant", **delta} + chunk = {"id": completion_id, "object": "chat.completion.chunk", "created": created, "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": None}]} + return f"data: {json.dumps(chunk, separators=(',', ':'))}\n\n" + + +def _could_be_tool_protocol(text: str) -> bool: + stripped = text.lstrip() + return _TOOL_PREFIX.startswith(stripped) or stripped.startswith(_TOOL_PREFIX) + + +def stream_chat_completion(engine: ForgeEngine, payload: Mapping[str, Any]) -> Iterable[str]: + messages, tools, max_tokens, temperature, session_id, stop = _request(payload) + completion_id, created = f"chatcmpl-{uuid.uuid4().hex}", int(time.time()) + model = str(payload.get("model") or engine.backend.name) + first, buffered, buffering_tool = True, "", False + for event in engine.stream_inference(messages, tools, max_tokens, temperature, session_id, stop): + if event.kind == "token": + if first or buffering_tool: + buffered += event.text + if _could_be_tool_protocol(buffered): + buffering_tool = True + continue + if buffered: + yield _content_chunk(completion_id, created, model, buffered, first) + first, buffered, buffering_tool = False, "", False + else: + yield _content_chunk(completion_id, created, model, event.text, False) + elif event.kind == "tool_call" and event.tool_call is not None: + # The native worker already validated the buffered text. Standard clients + # receive only the structured delta, never Forge's internal text protocol. + buffered, buffering_tool = "", False + delta = {"tool_calls": [{"index": 0, **_tool_call(event.tool_call)}]} + if first: delta = {"role": "assistant", **delta}; first = False + chunk = {"id": completion_id, "object": "chat.completion.chunk", "created": created, "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": None}]} + yield f"data: {json.dumps(chunk, separators=(',', ':'))}\n\n" + elif event.kind == "done": + if buffered: + yield _content_chunk(completion_id, created, model, buffered, first) + first, buffered = False, "" + chunk = {"id": completion_id, "object": "chat.completion.chunk", "created": created, "model": model, + "choices": [{"index": 0, "delta": {}, "finish_reason": _finish_reason(event.finish_reason)}], + "forge_session_id": session_id} + yield f"data: {json.dumps(chunk, separators=(',', ':'))}\n\n" + yield "data: [DONE]\n\n" diff --git a/extra/radeon_forge/ui/server.py b/extra/radeon_forge/ui/server.py new file mode 100644 index 0000000000000..ccde88679dab2 --- /dev/null +++ b/extra/radeon_forge/ui/server.py @@ -0,0 +1,183 @@ +from __future__ import annotations + +import argparse, json, mimetypes, re, sys, time +from dataclasses import asdict +from http import HTTPStatus +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Any, Iterable +from urllib.parse import parse_qs, urlparse + +from ..backends.fake import ScriptedBackend +from ..runtime import ForgeEngine, JsonlProcessBackend +from .openai_api import OpenAIRequestError, collect_chat_completion, stream_chat_completion + + +class ForgeRequestHandler(BaseHTTPRequestHandler): + engine: ForgeEngine + static_root = Path(__file__).parent / "static" + + def log_message(self, format, *args): print(f"[forge-ui] {self.address_string()} {format % args}") + + def _json(self, payload: Any, status: int = 200): + body = json.dumps(payload, default=str).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.send_header("Cache-Control", "no-store") + self.end_headers() + self.wfile.write(body) + + def _sse(self, chunks: Iterable[str]): + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Cache-Control", "no-cache") + self.send_header("Connection", "close") + self.end_headers() + try: + for chunk in chunks: + self.wfile.write(chunk.encode()) + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): pass + finally: self.close_connection = True + + def _body(self) -> dict[str, Any]: + length = int(self.headers.get("Content-Length", "0")) + payload = json.loads(self.rfile.read(length) or b"{}") + if not isinstance(payload, dict): raise ValueError("request body must be a JSON object") + return payload + + def _static(self, name: str): + path = (self.static_root / name).resolve() + if self.static_root.resolve() not in path.parents and path != self.static_root.resolve(): return self.send_error(404) + if not path.is_file(): return self.send_error(404) + body = path.read_bytes() + self.send_response(200) + self.send_header("Content-Type", mimetypes.guess_type(path.name)[0] or "application/octet-stream") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def _session_route(self): + match = re.fullmatch(r"/api/sessions/([a-f0-9]+)(?:/(.*))?", urlparse(self.path).path) + if not match: return None, None + return self.engine.session(match.group(1)), match.group(2) or "" + + def do_GET(self): + parsed, path = urlparse(self.path), urlparse(self.path).path + try: + if path == "/": return self._static("index.html") + if path in {"/app.js", "/style.css", "/kernels.css"}: return self._static(path[1:]) + if path == "/v1/models": + return self._json({"object":"list", "data":[{"id":self.engine.backend.name, "object":"model", + "created":int(time.time()), "owned_by":"local-radeon-forge", + "forge": {"runtime_fingerprint":asdict(self.engine.runtime_fingerprint()), + "capabilities":asdict(self.engine.backend.capabilities)}}]}) + if path == "/api/health": return self._json({"ok": True, "backend": self.engine.backend.name, + "capabilities": asdict(self.engine.backend.capabilities), "runtime_fingerprint": asdict(self.engine.runtime_fingerprint())}) + if path == "/api/sessions": return self._json(self.engine.sessions()) + if path == "/api/optimization": return self._json(self.engine.optimization_state()) + job_match = re.fullmatch(r"/api/jobs/([a-f0-9]+)", path) + if job_match: return self._json(asdict(self.engine.jobs.snapshot(job_match.group(1)))) + session, tail = self._session_route() + if session is None: return self.send_error(404) + if tail == "": return self._json(session.snapshot()) + if tail == "events": + after = int(parse_qs(parsed.query).get("after", ["0"])[0]) + return self._json({"events": session.events_after(after), "session": session.snapshot()}) + if tail == "profile": return self._json(self.engine.profile(session.session_id)) + if tail == "trace": return self._json(session.trace.to_dict()) + if tail == "trace/chrome": return self._json(session.trace.chrome_trace()) + return self.send_error(404) + except KeyError as exc: return self._json({"error": str(exc)}, 404) + except Exception as exc: return self._json({"error": str(exc), "type": type(exc).__name__}, 500) + + def do_POST(self): + path = urlparse(self.path).path + try: + body = self._body() + if path == "/v1/chat/completions": + if bool(body.get("stream", False)): return self._sse(stream_chat_completion(self.engine, body)) + return self._json(collect_chat_completion(self.engine, body)) + if path == "/api/sessions": return self._json(self.engine.create_session().snapshot(), HTTPStatus.CREATED) + if path == "/api/recipes/import": + result = self.engine.execute_explicit_ui_tool("import_forge_recipe", {"path": str(body.get("path", ""))}, + str(body.get("reason", "Import portable optimization recipe from local UI"))) + return self._json(result, HTTPStatus.CREATED) + if path == "/api/recipes/export": + result = self.engine.execute_explicit_ui_tool("export_forge_recipe", {"spec_id": str(body.get("spec_id", "")), + "candidate_id": body.get("candidate_id"), "output": str(body.get("output", ""))}, + str(body.get("reason", "Export portable execution-stage optimization recipe"))) + return self._json(result, HTTPStatus.CREATED) + if path == "/api/hooks/activate": + result = self.engine.execute_explicit_ui_tool("activate_optimization_hook", { + "candidate_id": str(body.get("candidate_id", "")), "reason": str(body.get("reason", "Deploy validated stage-specific optimization")), + "allow_unknown_runtime": bool(body.get("allow_unknown_runtime", False))}, + str(body.get("reason", "Deploy validated stage-specific optimization"))) + return self._json(result, HTTPStatus.CREATED) + if path == "/api/hooks/deactivate": + result = self.engine.execute_explicit_ui_tool("deactivate_optimization_hook", { + "activation_id": str(body.get("activation_id", "")), "reason": str(body.get("reason", "Rollback local optimization hook"))}, + str(body.get("reason", "Rollback local optimization hook"))) + return self._json(result) + session, tail = self._session_route() + if session is None: return self.send_error(404) + if tail == "messages": + events = session.send(str(body.get("content", "")), int(body.get("max_tokens", 512)), float(body.get("temperature", 0.0))) + return self._json({"events": [asdict(x) for x in events], "session": session.snapshot()}) + if tail == "messages/async": + job = self.engine.submit_message(session.session_id, str(body.get("content", "")), int(body.get("max_tokens", 512)), + float(body.get("temperature", 0.0))) + return self._json({"job": job, "session": session.snapshot()}, HTTPStatus.ACCEPTED) + if tail == "tools/approve": + reason = str(body.get("reason", "Approved from Radeon Forge UI")) + token = self.engine.grant_for_pending_tool(session.session_id, reason) + events = session.approve_tool(token, int(body.get("max_tokens", 512))) + return self._json({"events": [asdict(x) for x in events], "session": session.snapshot()}) + if tail == "tools/approve/async": + job = self.engine.submit_tool_approval(session.session_id, str(body.get("reason", "Approved from Radeon Forge UI")), + int(body.get("max_tokens", 512))) + return self._json({"job": job, "session": session.snapshot()}, HTTPStatus.ACCEPTED) + if tail == "tools/reject": + event = session.reject_tool(str(body.get("reason", "Rejected by user"))) + return self._json({"event": asdict(event), "session": session.snapshot()}) + return self.send_error(404) + except OpenAIRequestError as exc: + return self._json({"error":{"message":str(exc), "type":"invalid_request_error", "param":None, "code":None}}, 400) + except (ValueError, RuntimeError, FileNotFoundError) as exc: return self._json({"error": str(exc), "type": type(exc).__name__}, 400) + except Exception as exc: return self._json({"error": str(exc), "type": type(exc).__name__}, 500) + + +def build_backend(args): + if args.backend == "scripted": return ScriptedBackend() + if args.model is None: raise SystemExit("--model is required for --backend tinygrad-llama") + command = [sys.executable, "-m", "extra.radeon_forge.backends.tinygrad_llama_worker", "--model", str(args.model), + "--size", args.size, "--max-context", str(args.max_context)] + if args.tokenizer: command += ["--tokenizer", str(args.tokenizer)] + if args.quantize: command += ["--quantize", args.quantize] + return JsonlProcessBackend(command) + + +def main(): + parser = argparse.ArgumentParser(description="Radeon Forge local inference, agent and profiling server") + parser.add_argument("--workspace", type=Path, default=Path.cwd()) + parser.add_argument("--backend", choices=("scripted", "tinygrad-llama"), default="scripted") + parser.add_argument("--model", type=Path) + parser.add_argument("--tokenizer", type=Path) + parser.add_argument("--size", choices=("1B", "8B", "70B", "405B"), default="1B") + parser.add_argument("--quantize", choices=("int8", "nf4", "float16", "fp8")) + parser.add_argument("--max-context", type=int, default=8192) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=7790) + args = parser.parse_args() + if args.host not in {"127.0.0.1", "localhost", "::1"}: raise SystemExit("Forge binds to loopback only") + engine = ForgeEngine(build_backend(args), args.workspace) + ForgeRequestHandler.engine = engine + server = ThreadingHTTPServer((args.host, args.port), ForgeRequestHandler) + print(f"Radeon Forge: http://{args.host}:{args.port} backend={engine.backend.name} workspace={args.workspace.resolve()}") + print(f"OpenAI-compatible API: http://{args.host}:{args.port}/v1") + try: server.serve_forever() + finally: engine.close() + + +if __name__ == "__main__": main() diff --git a/extra/radeon_forge/ui/static/app.js b/extra/radeon_forge/ui/static/app.js new file mode 100644 index 0000000000000..166fcf6feb609 --- /dev/null +++ b/extra/radeon_forge/ui/static/app.js @@ -0,0 +1,33 @@ +const state={session:null,sessions:[],health:null,optimization:null,tab:'console',activeJob:null}; +const $=s=>document.querySelector(s), $$=s=>[...document.querySelectorAll(s)], sleep=ms=>new Promise(r=>setTimeout(r,ms)); +const toast=msg=>{const el=$('#toast');el.textContent=msg;el.classList.add('show');setTimeout(()=>el.classList.remove('show'),3000)}; +async function api(path,opts={}){const res=await fetch(path,{headers:{'Content-Type':'application/json'},...opts});const data=await res.json();if(!res.ok)throw new Error(data.error||`HTTP ${res.status}`);return data} +function fmtMs(v){return v==null?'—':`${Number(v).toFixed(v<10?2:1)} ms`} +function esc(v){return String(v??'').replace(/[&<>"']/g,m=>({'&':'&','<':'<','>':'>','"':'"',"'":'''}[m]))} +function hookWhen(h){const w=h?.when||{};const stages=(w.stages||['any']).map(String).join(', ');const detail=[];if(w.min_context_tokens!=null)detail.push(`ctx≥${w.min_context_tokens}`);if(w.max_context_tokens!=null)detail.push(`ctx≤${w.max_context_tokens}`);if(w.min_generated_token_index!=null)detail.push(`token≥${w.min_generated_token_index}`);if(w.max_generated_token_index!=null)detail.push(`token≤${w.max_generated_token_index}`);if(w.batch_sizes?.length)detail.push(`batch=${w.batch_sizes.join('/')}`);if(w.prefix_cache&&w.prefix_cache!=='any')detail.push(`cache=${w.prefix_cache}`);if(w.warm!=null)detail.push(w.warm?'warm':'cold');Object.entries(w.conditions||{}).forEach(([k,v])=>detail.push(`${k}=${v}`));return {stages,detail:detail.join(' · ')}} +async function boot(){try{state.health=await api('/api/health');renderHealth();await refreshSessions();if(!state.sessions.length)await createSession();else await selectSession(state.sessions[0].session_id)}catch(e){toast(e.message)}} +function renderHealth(){const h=state.health;$('#runtime-name').textContent=h?.backend||'Offline';const fp=h?.runtime_fingerprint||{};$('#backend-status').textContent=h?`${h.backend} · ${fp.architecture||'unbound architecture'} · local loopback runtime`:'Backend unavailable';$('#runtime-arch').textContent=fp.architecture||'unbound';$('#runtime-model').textContent=fp.model_family||'—';$('#runtime-dtype').textContent=fp.dtype||'—';const caps=$('#capabilities');caps.innerHTML='';Object.entries(h?.capabilities||{}).forEach(([k,v])=>{const el=document.createElement('span');el.className=`cap ${v?'on':''}`;el.textContent=k.replaceAll('_',' ');caps.append(el)})} +async function refreshSessions(){state.sessions=await api('/api/sessions');renderSessionList()} +function renderSessionList(){const list=$('#session-list');list.innerHTML='';state.sessions.forEach((s,i)=>{const el=document.createElement('div');el.className=`session-item ${state.session?.session_id===s.session_id?'active':''}`;el.innerHTML=`${i===0?'Primary':'Session'} ${s.session_id.slice(0,6)}
${esc(s.state)}`;el.onclick=()=>selectSession(s.session_id);list.append(el)})} +async function createSession(){const s=await api('/api/sessions',{method:'POST',body:'{}'});await refreshSessions();await selectSession(s.session_id)} +async function selectSession(id){if(state.activeJob)throw new Error('A local turn is still running');state.session=await api(`/api/sessions/${id}`);renderSession();renderSessionList();if(state.tab==='profile')refreshProfile();if(state.tab==='trace')refreshTrace()} +function renderSession(){const s=state.session;if(!s)return;$('#session-title').textContent=`Agent Console · ${s.session_id.slice(0,6)}`;$('#session-state').textContent=s.state;$('#trace-id').textContent=s.trace_id.slice(0,10);$('#event-count').textContent=s.events.length;renderMessages();renderApproval()} +function renderMessages(){const box=$('#messages');const messages=(state.session?.messages||[]).filter(x=>x.role!=='system'&&x.role!=='tool');if(!messages.length&&!state.session?.partial_output)return;box.innerHTML='';messages.forEach(m=>{const el=document.createElement('div');el.className=`message ${m.role}`;el.innerHTML=`
${m.role==='user'?'YOU':'RF'}
${esc(m.content)}
`;box.append(el)});if(state.session?.partial_output){const el=document.createElement('div');el.className='message assistant streaming';el.innerHTML=`
RF
${esc(state.session.partial_output)}
`;box.append(el)}box.scrollTop=box.scrollHeight} +function renderApproval(){const card=$('#approval-card'),call=state.session?.pending_tool_call;if(!call){card.classList.add('hidden');return}card.classList.remove('hidden');$('#approval-name').textContent=call.name;$('#approval-args').textContent=JSON.stringify(call.arguments,null,2)} +function maxSequence(){return Math.max(0,...(state.session?.events||[]).map(x=>Number(x.sequence)||0))} +async function watchJob(jobId){state.activeJob=jobId;let after=maxSequence(),profileTick=0;try{while(true){const [job,delta]=await Promise.all([api(`/api/jobs/${jobId}`),api(`/api/sessions/${state.session.session_id}/events?after=${after}`)]);state.session=delta.session;(delta.events||[]).forEach(e=>after=Math.max(after,Number(e.sequence)||0));renderSession();if(++profileTick%5===0&&state.tab==='trace')refreshTrace().catch(()=>{});if(job.state==='completed')break;if(job.state==='failed')throw new Error(`${job.error_type||'JobError'}: ${job.error||'local job failed'}`);await sleep(100)}}finally{state.activeJob=null;await refreshSessions();state.session=await api(`/api/sessions/${state.session.session_id}`);renderSession();if(state.tab==='profile')await refreshProfile();if(state.tab==='trace')await refreshTrace();if(state.tab==='kernels')await refreshKernels()}} +async function sendMessage(text){if(!state.session||state.activeJob)return;$('#send').disabled=true;$('#send').textContent='Running…';try{const out=await api(`/api/sessions/${state.session.session_id}/messages/async`,{method:'POST',body:JSON.stringify({content:text,max_tokens:512})});state.session=out.session;renderSession();await watchJob(out.job.job_id)}finally{$('#send').disabled=false;$('#send').textContent='Run'}} +async function approveTool(){if(state.activeJob)return;$('#approve-tool').disabled=true;try{const out=await api(`/api/sessions/${state.session.session_id}/tools/approve/async`,{method:'POST',body:JSON.stringify({reason:'Approved once from the local UI',max_tokens:512})});state.session=out.session;renderSession();await watchJob(out.job.job_id);toast('Tool executed locally')}finally{$('#approve-tool').disabled=false}} +async function rejectTool(){const out=await api(`/api/sessions/${state.session.session_id}/tools/reject`,{method:'POST',body:JSON.stringify({reason:'Rejected in UI'})});state.session=out.session;renderSession();toast('Tool rejected')} +async function refreshProfile(){if(!state.session)return;const p=await api(`/api/sessions/${state.session.session_id}/profile`),s=p.summary||{},stages=s.token_latency_by_stage||{};$('#p50-token').textContent=fmtMs(s.token_wall_ms_p50);$('#p95-token').textContent=fmtMs(s.token_wall_ms_p95);$('#first-token').textContent=fmtMs(stages.first_token?.p50_ms);$('#steady-decode').textContent=fmtMs(stages.decode?.p50_ms);$('#launches-token').textContent=s.mean_kernel_count_per_token==null?'—':Number(s.mean_kernel_count_per_token).toFixed(1);$('#hook-transitions').textContent=`${s.hook_transitions||0}${s.hook_rollbacks?` / ${s.hook_rollbacks} rollback`:''}`;const findings=$('#findings');findings.innerHTML='';(p.findings||[]).forEach(f=>{const el=document.createElement('div');el.className=`finding ${f.severity}`;el.innerHTML=`
${esc(f.status)}${esc(f.title)}

${esc((f.evidence||[]).join(' · '))}

Next: ${esc(f.recommendation)}

`;findings.append(el)});if(!p.findings?.length)findings.innerHTML='

No diagnosis yet. Run an agent turn first.

';renderAttribution(s.durations_ms_by_kind||{});renderStageKernels(s.top_kernels_by_stage||{})} +function renderAttribution(d){const box=$('#attribution');box.innerHTML='';const entries=Object.entries(d).sort((a,b)=>b[1]-a[1]),max=Math.max(1,...entries.map(x=>x[1]));entries.forEach(([k,v])=>{const el=document.createElement('div');el.className='bar-row';el.innerHTML=`
${esc(k)}${fmtMs(v)}
`;box.append(el)});if(!entries.length)box.textContent='No trace data yet.'} +function renderStageKernels(byStage){const box=$('#stage-kernels');box.innerHTML='';Object.entries(byStage).sort().forEach(([stage,rows])=>{const top=rows?.[0];if(!top)return;const el=document.createElement('div');el.className='kernel-card compact-card';el.innerHTML=`
${esc(stage)}${esc(top.name)}

${top.calls} calls · ${fmtMs(top.total_ms)} · ${(100*top.share).toFixed(1)}% of captured stage time

`;box.append(el)});if(!box.children.length)box.innerHTML='

No stage-level kernel evidence yet.

'} +async function refreshTrace(){if(!state.session)return;const t=await api(`/api/sessions/${state.session.session_id}/trace`),box=$('#timeline');box.innerHTML='';(t.events||[]).forEach(e=>{const el=document.createElement('div');el.className=`trace-row ${e.name==='hook'?'hook-row':''}`;const attrs=Object.entries(e.attributes||{}).slice(0,5).map(([k,v])=>`${k}=${Array.isArray(v)?v.join(','):v}`).join(' · ');el.innerHTML=`${esc(e.kind)}${esc(e.name)}${esc(attrs)}${fmtMs(e.duration_ms)}`;box.append(el)});if(!t.events?.length)box.innerHTML='

No timeline events yet.

'} +async function exportCandidate(candidate){const output=`exports/${candidate.candidate_id}.forge.toml`;const out=await api('/api/recipes/export',{method:'POST',body:JSON.stringify({spec_id:candidate.spec_id,candidate_id:candidate.candidate_id,output,reason:'Export shareable stage-aware optimization recipe'})});toast(`Exported ${out.path}`);await refreshKernels()} +async function importRecipe(){const input=$('#recipe-path'),path=input.value.trim();if(!path)return;const out=await api('/api/recipes/import',{method:'POST',body:JSON.stringify({path,reason:'Import local portable optimization recipe'})});input.value='';toast(`Installed ${out.recipe_id}`);await refreshKernels()} +async function activateCandidate(candidate){const reason=`Activate ${candidate.candidate_id} only in its validated execution states`;const out=await api('/api/hooks/activate',{method:'POST',body:JSON.stringify({candidate_id:candidate.candidate_id,reason})});toast(`Activated ${out.activation_id.slice(0,8)}`);await refreshKernels()} +async function rollbackHook(hook){const out=await api('/api/hooks/deactivate',{method:'POST',body:JSON.stringify({activation_id:hook.activation_id,reason:`Rollback ${hook.candidate_id} from local UI`})});toast(`Rolled back ${out.candidate_id}`);await refreshKernels()} +function renderActiveHooks(hooks){const box=$('#active-hooks');box.innerHTML='';(hooks||[]).forEach(h=>{const w=hookWhen(h.descriptor),el=document.createElement('div');el.className='kernel-card active-hook';el.innerHTML=`
active${esc(h.descriptor.layer)} · ${esc(h.descriptor.target)}

${esc(h.candidate_id)}

${esc(w.stages)}${esc(h.descriptor.mode)}priority ${h.descriptor.priority||0}
${w.detail?`

${esc(w.detail)}

`:''}
${esc(JSON.stringify({selector:h.descriptor.selector,compatibility:h.compatibility},null,2))}
`;el.querySelector('.rollback-hook').onclick=()=>rollbackHook(h).catch(e=>toast(e.message));box.append(el)});if(!box.children.length)box.innerHTML='

No optimization is active. The trusted tinygrad baseline owns every execution stage.

'} +async function refreshKernels(){state.optimization=await api('/api/optimization');const specs=$('#kernel-specs'),candidates=$('#kernel-candidates'),recipes=$('#recipe-list');specs.innerHTML='';candidates.innerHTML='';recipes.innerHTML='';renderActiveHooks(state.optimization.active_hooks||[]);const activeIds=new Set((state.optimization.active_hooks||[]).map(x=>x.candidate_id));(state.optimization.specs||[]).forEach(s=>{const h=s.metadata?.hook||{},w=hookWhen(h),el=document.createElement('div');el.className='kernel-card';el.innerHTML=`

${esc(s.name)}

${esc(s.operation)}

${esc(s.target)}${esc(h.layer||'kernel')}${esc(w.stages)}${esc(s.objective)}
${w.detail?`

${esc(w.detail)}

`:''}
${esc(JSON.stringify({where:{layer:h.layer,target:h.target,mode:h.mode,selector:h.selector},when:h.when,shapes:s.shapes,dtypes:s.dtypes,invariants:s.invariants,agent_brief:s.metadata?.recipe_agent_brief},null,2))}
`;specs.append(el)});(state.optimization.candidates||[]).forEach(c=>{const spec=(state.optimization.specs||[]).find(s=>s.spec_id===c.spec_id),h=spec?.metadata?.hook||{},w=hookWhen(h),eligible=['hardware_passed','heldout_passed','accepted'].includes(c.status),el=document.createElement('div');el.className='kernel-card';el.innerHTML=`

${esc(c.candidate_id)}

${esc(c.hypothesis||'No hypothesis recorded')}

${esc(c.status)}${esc(h.layer||'kernel')}${esc(w.stages)}
${esc(JSON.stringify(c.evidence||{},null,2))}
${eligible&&!activeIds.has(c.candidate_id)?'':''}${activeIds.has(c.candidate_id)?'currently active':''}
`;el.querySelector('.export-candidate').onclick=()=>exportCandidate(c).catch(e=>toast(e.message));const activate=el.querySelector('.activate-candidate');if(activate)activate.onclick=()=>activateCandidate(c).catch(e=>toast(e.message));candidates.append(el)});(state.optimization.recipes||[]).forEach(r=>{const el=document.createElement('div');el.className='kernel-card';el.innerHTML=`

${esc(r.recipe_id)}

Installed portable contract. Placement and execution-state predicates are preserved; cached code remains unverified until local oracles pass.

${esc(r.spec_id)}${r.seed_candidate_id?'seed cache':''}
${esc(JSON.stringify({bundle:r.bundle_root,artifacts:r.artifact_hashes},null,2))}
`;recipes.append(el)});if(!state.optimization.candidates?.length)candidates.innerHTML='

No candidate has been staged. Ask Forge to list contracts and generate one for a specific execution state.

';if(!state.optimization.recipes?.length)recipes.innerHTML='

No portable recipe installed yet. Import a .forge.toml file from the workspace.

'} +function switchTab(tab){state.tab=tab;$$('.tab').forEach(x=>x.classList.remove('active'));$(`#${tab}-tab`).classList.add('active');$$('[data-tab]').forEach(x=>x.classList.toggle('active',x.dataset.tab===tab));if(tab==='profile')refreshProfile();if(tab==='trace')refreshTrace();if(tab==='kernels')refreshKernels()} +$('#new-session').onclick=()=>createSession().catch(e=>toast(e.message));$('#composer').onsubmit=async e=>{e.preventDefault();const input=$('#prompt'),text=input.value.trim();if(!text)return;input.value='';try{await sendMessage(text)}catch(err){toast(err.message)}};$('#approve-tool').onclick=()=>approveTool().catch(e=>toast(e.message));$('#reject-tool').onclick=()=>rejectTool().catch(e=>toast(e.message));$('#refresh-profile').onclick=()=>refreshProfile().catch(e=>toast(e.message));$('#refresh-trace').onclick=()=>refreshTrace().catch(e=>toast(e.message));$('#refresh-kernels').onclick=()=>refreshKernels().catch(e=>toast(e.message));$('#import-recipe').onclick=()=>importRecipe().catch(e=>toast(e.message));$$('[data-tab]').forEach(x=>x.onclick=()=>switchTab(x.dataset.tab));boot(); diff --git a/extra/radeon_forge/ui/static/index.html b/extra/radeon_forge/ui/static/index.html new file mode 100644 index 0000000000000..c81ef3102a988 --- /dev/null +++ b/extra/radeon_forge/ui/static/index.html @@ -0,0 +1,95 @@ + + + + + + Radeon Forge + + + + +
+ + +
+
+

Agent Console

Connecting to local backend…

+
+ + + + +
+
+ +
+
+
RF

Optimize the whole agent loop.

Run a private model, inspect the repository, approve tools, and trace every turn down to GPU kernels and execution stages.

+
+ +
+ + +
+
+ +
+
+
P50 token—all generated tokens
+
P95 token—all generated tokens
+
First token—P50 execution state
+
Steady decode—P50 execution state
+
Launches/token—tinygrad counters
+
Hook transitions0including rollbacks
+
+
+

Agentic diagnosis

+

Time attribution

+

Dominant kernels by execution state

+
+
+ +
+

Unified execution timeline

Agent → model state → token → tool → kernel → optimization switch

+
+ +
+
+

Active execution hooks

Where the optimization intercepts × when it is allowed to run.

+

Typed optimization contracts

Intent and invariants are durable; generated implementations are disposable.

+

Candidate pipeline

Staged → MockGPU → W7900 → held-out → active / rejected

+

Portable optimization recipes

Share placement, execution-state predicate, intent, oracle, knowledge, and an optional implementation cache in one file.

+
+
+
+
+
+
+ + +
+
+ + + diff --git a/extra/radeon_forge/ui/static/kernels.css b/extra/radeon_forge/ui/static/kernels.css new file mode 100644 index 0000000000000..e7bba3e605a2c --- /dev/null +++ b/extra/radeon_forge/ui/static/kernels.css @@ -0,0 +1 @@ +.metric-grid{grid-template-columns:repeat(3,minmax(0,1fr))}.profile-layout{overflow:auto}.kernel-layout{display:grid;grid-template-columns:1fr 1fr;gap:12px;padding:22px 24px 24px;width:100%;min-height:0;overflow:auto}.active-hooks-panel,.recipe-panel{grid-column:1/-1}.kernel-list{padding:12px;overflow:auto}.kernel-list.compact{display:grid;grid-template-columns:repeat(auto-fit,minmax(220px,1fr));gap:8px}.kernel-card{background:var(--panel2);border:1px solid var(--line);border-radius:12px;padding:14px;margin-bottom:10px}.kernel-card h4{margin:0 0 6px;font-size:14px}.kernel-card p{color:var(--muted);font-size:11px;line-height:1.5}.kernel-card.active-hook{border-color:rgba(61,220,151,.32);box-shadow:inset 3px 0 0 rgba(61,220,151,.65)}.compact-card{margin:0}.stage-line{display:flex;align-items:center;gap:8px;margin-bottom:7px}.stage-line strong{font-size:12px;overflow:hidden;text-overflow:ellipsis;white-space:nowrap}.kernel-meta{display:flex;flex-wrap:wrap;gap:6px;margin-top:10px}.status{padding:4px 7px;border-radius:7px;background:#262d3a;border:1px solid var(--line);font-size:9px;text-transform:uppercase;letter-spacing:.07em}.status.stage{color:#cbb7ff;border-color:rgba(167,67,255,.3);background:rgba(167,67,255,.1)}.status.active{color:var(--good);border-color:rgba(61,220,151,.3);background:rgba(61,220,151,.1)}.status.mockgpu_passed,.status.hardware_passed,.status.heldout_passed,.status.accepted{color:var(--good);border-color:rgba(61,220,151,.28);background:rgba(61,220,151,.08)}.status.mockgpu_failed,.status.hardware_failed,.status.heldout_failed{color:#ff8e98;border-color:rgba(255,95,109,.28);background:rgba(255,95,109,.08)}.status.imported_unverified{color:#ffd278;border-color:rgba(245,185,66,.28);background:rgba(245,185,66,.08)}.contract{margin-top:10px;padding:9px;background:#0f1219;border:1px solid var(--line);border-radius:9px;font:10px/1.5 ui-monospace,SFMono-Regular,Menlo,monospace;color:#aeb8c8;white-space:pre-wrap;max-height:170px;overflow:auto}.recipe-import{display:flex;gap:10px;padding:14px 14px 0}.recipe-import input{flex:1;background:#0f1219;border:1px solid var(--line);border-radius:10px;color:var(--text);padding:10px 12px;outline:0}.recipe-import input:focus{border-color:rgba(255,92,53,.65)}.kernel-actions{display:flex;align-items:center;flex-wrap:wrap;gap:8px;margin-top:10px}.kernel-actions button{font-size:10px}.hook-row{background:rgba(167,67,255,.06)}.stage-panel{grid-column:1/-1}@media(max-width:900px){.metric-grid{grid-template-columns:repeat(2,minmax(0,1fr))}.kernel-layout{grid-template-columns:1fr}.active-hooks-panel,.recipe-panel{grid-column:auto}} diff --git a/extra/radeon_forge/ui/static/style.css b/extra/radeon_forge/ui/static/style.css new file mode 100644 index 0000000000000..73a72cf69dab7 --- /dev/null +++ b/extra/radeon_forge/ui/static/style.css @@ -0,0 +1 @@ +:root{--bg:#0b0d12;--panel:#11151d;--panel2:#171c26;--line:#252c39;--text:#edf2f7;--muted:#8e99aa;--accent:#ff5c35;--accent2:#ff8b69;--good:#3ddc97;--warn:#f5b942;--bad:#ff5f6d;font-family:Inter,ui-sans-serif,system-ui,-apple-system,BlinkMacSystemFont,"Segoe UI",sans-serif;color-scheme:dark}*{box-sizing:border-box}body{margin:0;background:radial-gradient(circle at 50% -20%,#1b2130 0,#0b0d12 34%);color:var(--text);height:100vh;overflow:hidden}.shell{display:grid;grid-template-columns:248px minmax(0,1fr) 286px;height:100vh}.rail,.inspector{background:rgba(11,13,18,.92);backdrop-filter:blur(18px);padding:22px;border-color:var(--line)}.rail{border-right:1px solid var(--line);display:flex;flex-direction:column}.inspector{border-left:1px solid var(--line)}.brand{display:flex;gap:12px;align-items:center;margin-bottom:28px}.brand strong,.brand span{display:block}.brand span{font-size:12px;color:var(--muted);margin-top:2px}.mark,.hero-mark{display:grid;place-items:center;background:linear-gradient(135deg,var(--accent),#b32dff);box-shadow:0 10px 30px rgba(255,92,53,.22);font-weight:800}.mark{width:38px;height:38px;border-radius:12px;font-size:13px}.hero-mark{width:64px;height:64px;border-radius:20px;margin:auto;font-size:19px}.primary,.send,.ghost,.danger{border:0;border-radius:10px;font-weight:650;cursor:pointer;transition:.18s ease}.primary,.send{background:linear-gradient(135deg,var(--accent),var(--accent2));color:white;padding:11px 14px;box-shadow:0 8px 20px rgba(255,92,53,.17)}.primary:hover,.send:hover{transform:translateY(-1px);filter:brightness(1.05)}.ghost{background:transparent;color:var(--muted);padding:8px 10px;border:1px solid transparent}.ghost:hover,.ghost.active{color:var(--text);background:var(--panel2);border-color:var(--line)}.danger{background:rgba(255,95,109,.12);color:#ff8e98;padding:10px 13px;border:1px solid rgba(255,95,109,.25)}.rail-label,.eyebrow{text-transform:uppercase;letter-spacing:.12em;font-size:10px;color:var(--muted);font-weight:750}.rail-label{margin:26px 0 10px}.session-list{display:flex;flex-direction:column;gap:6px;overflow:auto}.session-item{padding:10px;border-radius:10px;color:var(--muted);cursor:pointer;border:1px solid transparent;font-size:13px}.session-item:hover,.session-item.active{background:var(--panel2);border-color:var(--line);color:var(--text)}.privacy-card{margin-top:auto;display:flex;gap:10px;padding:12px;background:var(--panel);border:1px solid var(--line);border-radius:13px}.privacy-card strong,.privacy-card small{display:block}.privacy-card strong{font-size:12px}.privacy-card small{font-size:10px;color:var(--muted);line-height:1.4;margin-top:3px}.dot{width:8px;height:8px;background:var(--good);border-radius:50%;box-shadow:0 0 0 5px rgba(61,220,151,.1);margin-top:4px}.main{min-width:0;display:flex;flex-direction:column}.topbar{height:82px;border-bottom:1px solid var(--line);padding:15px 24px;display:flex;align-items:center;justify-content:space-between;background:rgba(11,13,18,.55);backdrop-filter:blur(14px)}h1,h2,h3,p{margin:0}.topbar h1{font-size:18px}.topbar p{color:var(--muted);font-size:11px;margin-top:5px}.top-actions{display:flex;gap:4px}.tab{display:none;min-height:0;flex:1}.tab.active{display:flex}.messages{flex:1;overflow:auto;padding:28px max(28px,8vw) 160px}.empty-state{text-align:center;max-width:580px;margin:15vh auto 0}.empty-state h2{font-size:30px;margin:20px 0 10px}.empty-state p{color:var(--muted);line-height:1.7}.message{max-width:780px;margin:0 auto 18px;display:flex;gap:12px}.message .avatar{width:30px;height:30px;border-radius:9px;display:grid;place-items:center;flex:none;background:var(--panel2);font-size:10px;font-weight:800}.message.user .avatar{background:rgba(255,92,53,.18);color:var(--accent2)}.bubble{padding:13px 15px;border:1px solid var(--line);background:var(--panel);border-radius:4px 14px 14px 14px;line-height:1.55;font-size:14px;white-space:pre-wrap;overflow-wrap:anywhere}.message.user .bubble{background:#171b24}.composer{position:absolute;bottom:20px;left:calc(248px + 5vw);right:calc(286px + 5vw);display:flex;gap:10px;padding:10px;border-radius:16px;background:rgba(22,27,37,.96);border:1px solid #30394a;box-shadow:0 20px 60px rgba(0,0,0,.45)}.composer textarea{flex:1;resize:none;background:transparent;border:0;outline:0;color:var(--text);font:inherit;padding:9px;max-height:160px}.send{align-self:flex-end}.approval-card{position:absolute;bottom:112px;left:calc(248px + 8vw);right:calc(286px + 8vw);display:flex;justify-content:space-between;gap:18px;padding:16px;background:#1a1717;border:1px solid rgba(255,139,105,.34);border-radius:14px;box-shadow:0 20px 50px rgba(0,0,0,.42)}.approval-card.hidden{display:none}.approval-card strong{display:block;margin:6px 0}.approval-card pre{margin:0;color:#d7bdaf;max-height:110px;overflow:auto}.approval-actions{display:flex;gap:8px;align-items:flex-end}.metric-grid{display:grid;grid-template-columns:repeat(4,1fr);gap:12px;padding:22px 24px 0;width:100%}.metric-grid article{background:var(--panel);border:1px solid var(--line);border-radius:14px;padding:16px}.metric-grid span,.metric-grid small{display:block;color:var(--muted);font-size:11px}.metric-grid strong{display:block;font-size:25px;margin:8px 0 4px}.profile-layout{display:grid;grid-template-columns:1.4fr .8fr;gap:12px;padding:12px 24px 24px;width:100%;overflow:hidden}.panel{background:var(--panel);border:1px solid var(--line);border-radius:14px;min-height:0;overflow:auto}.panel-head{display:flex;justify-content:space-between;align-items:center;padding:16px 18px;border-bottom:1px solid var(--line)}.panel-head p{font-size:11px;color:var(--muted);margin-top:4px}.findings{padding:10px}.finding{padding:14px;border-radius:11px;margin-bottom:8px;background:var(--panel2);border:1px solid var(--line)}.finding-head{display:flex;align-items:center;gap:8px}.badge{font-size:9px;text-transform:uppercase;letter-spacing:.08em;padding:4px 6px;border-radius:6px;background:#292f3c;color:var(--muted)}.finding.high .badge{background:rgba(255,95,109,.14);color:#ff8e98}.finding.medium .badge{background:rgba(245,185,66,.14);color:#ffd278}.finding p{color:var(--muted);font-size:12px;line-height:1.5;margin-top:9px}.attribution{padding:18px}.bar-row{margin-bottom:16px}.bar-label{display:flex;justify-content:space-between;color:var(--muted);font-size:11px;margin-bottom:7px}.bar{height:8px;background:#222936;border-radius:20px;overflow:hidden}.bar>i{display:block;height:100%;background:linear-gradient(90deg,var(--accent),#a743ff);border-radius:20px}.trace-panel{margin:22px 24px 24px;width:100%}.timeline{padding:16px;overflow:auto}.trace-row{display:grid;grid-template-columns:95px 145px minmax(120px,1fr) 80px;gap:10px;align-items:center;padding:9px 8px;border-bottom:1px solid rgba(37,44,57,.65);font-size:11px}.trace-kind{color:var(--accent2);text-transform:uppercase;letter-spacing:.07em}.trace-name{font-weight:650}.trace-attrs{color:var(--muted);white-space:nowrap;overflow:hidden;text-overflow:ellipsis}.trace-duration{text-align:right;color:#b8c1cf}.inspector-head{padding-bottom:18px}.inspector-head strong{display:block;font-size:16px;margin-top:5px}.caps{display:flex;flex-wrap:wrap;gap:6px}.cap{padding:5px 7px;border-radius:7px;background:var(--panel2);border:1px solid var(--line);font-size:10px;color:var(--muted)}.cap.on{color:var(--good);border-color:rgba(61,220,151,.25)}.divider{height:1px;background:var(--line);margin:20px 0}.facts{display:grid;grid-template-columns:70px 1fr;gap:10px;font-size:11px}.facts dt{color:var(--muted)}.facts dd{margin:0;text-align:right;white-space:nowrap;overflow:hidden;text-overflow:ellipsis}.flow{margin:12px 0 0;padding-left:20px;color:var(--muted);font-size:11px;line-height:2}.flow li::marker{color:var(--accent)}.toast{position:fixed;right:20px;bottom:20px;padding:11px 14px;background:#222834;border:1px solid #394253;border-radius:10px;opacity:0;transform:translateY(8px);pointer-events:none;transition:.2s}.toast.show{opacity:1;transform:none}@media(max-width:1100px){.shell{grid-template-columns:210px 1fr}.inspector{display:none}.composer{left:calc(210px + 4vw);right:4vw}.approval-card{left:calc(210px + 6vw);right:6vw}}@media(max-width:760px){.shell{display:block}.rail{display:none}.main{height:100vh}.top-actions{display:none}.composer{left:14px;right:14px}.approval-card{left:14px;right:14px}.metric-grid{grid-template-columns:1fr 1fr}.profile-layout{grid-template-columns:1fr}.trace-row{grid-template-columns:70px 110px 1fr}.trace-duration{display:none}} \ No newline at end of file diff --git a/extra/radeon_forge/workloads/__init__.py b/extra/radeon_forge/workloads/__init__.py new file mode 100644 index 0000000000000..15c1fb0f09635 --- /dev/null +++ b/extra/radeon_forge/workloads/__init__.py @@ -0,0 +1 @@ +"""Hardware workload adapters for Radeon Forge.""" diff --git a/extra/radeon_forge/workloads/block_hook_harness.py b/extra/radeon_forge/workloads/block_hook_harness.py new file mode 100644 index 0000000000000..343a2642dce55 --- /dev/null +++ b/extra/radeon_forge/workloads/block_hook_harness.py @@ -0,0 +1,192 @@ +from __future__ import annotations + +import argparse, hashlib, importlib.util, json, os, time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping + +import numpy as np + +from tinygrad import Device, GlobalCounters, Tensor, dtypes +from extra.models.llama import TransformerBlock, precompute_freqs_cis + + +@dataclass(frozen=True) +class HarnessConfig: + dim: int = 128 + hidden_dim: int = 256 + n_heads: int = 4 + n_kv_heads: int = 4 + max_context: int = 128 + prompt_tokens: int = 8 + dtype: str = "float32" + max_abs_error: float = 2e-4 + max_rel_error: float = 2e-3 + + +class ModelShell: + def __init__(self, layer: Any, config: HarnessConfig): + self.layers = [layer] + self.max_context = config.max_context + self.forward_jit = None + + +def _candidate_path(argument: str | None) -> Path: + value = argument or os.environ.get("RADEON_FORGE_CANDIDATE") + if not value: raise ValueError("candidate path is required") + path = Path(value).resolve() + if not path.is_file(): raise FileNotFoundError(path) + return path + + +def _parameters() -> dict[str, Any]: + raw = os.environ.get("RADEON_FORGE_CANDIDATE_JSON", "") + if not raw: return {} + payload = json.loads(raw) + values = payload.get("parameters", {}) if isinstance(payload, Mapping) else {} + if not isinstance(values, Mapping): raise ValueError("candidate parameters must be an object") + return dict(values) + + +def _load_candidate(path: Path, parameters: Mapping[str, Any]): + digest = hashlib.sha256(path.read_bytes()).hexdigest() + spec = importlib.util.spec_from_file_location(f"radeon_forge_candidate_{digest[:16]}", path) + if spec is None or spec.loader is None: raise RuntimeError(f"cannot import candidate {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + setattr(module, "RADEON_FORGE_PARAMETERS", dict(parameters)) + configure = getattr(module, "configure", None) + if configure is not None: + if not callable(configure): raise TypeError("optional configure must be callable") + configure(dict(parameters)) + factory = getattr(module, "build_replacement", None) + if not callable(factory): raise TypeError("candidate must define build_replacement(original, layer_index, model, context)") + return factory + + +def _block(config: HarnessConfig) -> TransformerBlock: + return TransformerBlock(config.dim, config.hidden_dim, config.n_heads, config.n_kv_heads, + norm_eps=1e-5, max_context=config.max_context) + + +def _reset_cache(block: TransformerBlock) -> None: + attention = getattr(block, "attention", None) + if attention is not None and hasattr(attention, "cache_kv"): delattr(attention, "cache_kv") + + +def _dtype(name: str): + return {"float32":dtypes.float32, "float16":dtypes.float16, "bfloat16":dtypes.bfloat16}[name] + + +def _inputs(config: HarnessConfig, prompt_tokens: int, seed: int): + Tensor.manual_seed(seed) + dtype = _dtype(config.dtype) + prompt = Tensor.randn(1, prompt_tokens, config.dim).cast(dtype).contiguous().realize() + first = Tensor.randn(1, 1, config.dim).cast(dtype).contiguous().realize() + decode = Tensor.randn(1, 1, config.dim).cast(dtype).contiguous().realize() + freqs = precompute_freqs_cis(config.dim // config.n_heads, config.max_context * 2).cast(dtype).contiguous().realize() + mask = Tensor.full((1, 1, prompt_tokens, prompt_tokens), float("-inf"), dtype=dtype).triu(1) + return prompt, first, decode, freqs, mask + + +def _run_transition(callable_block, config: HarnessConfig, prompt_tokens: int, seed: int) -> list[np.ndarray]: + prompt, first, decode, freqs, mask = _inputs(config, prompt_tokens, seed) + outputs = [ + callable_block(prompt, 0, freqs[:, :prompt_tokens], mask).realize().numpy(), + callable_block(first, prompt_tokens, freqs[:, prompt_tokens:prompt_tokens+1], None).realize().numpy(), + callable_block(decode, prompt_tokens+1, freqs[:, prompt_tokens+1:prompt_tokens+2], None).realize().numpy(), + ] + return outputs + + +def _error(reference: list[np.ndarray], candidate: list[np.ndarray]) -> dict[str, Any]: + max_abs = max(float(np.max(np.abs(left.astype(np.float64)-right.astype(np.float64)))) for left, right in zip(reference, candidate)) + max_rel = 0.0 + checked = 0 + for left, right in zip(reference, candidate): + left64, right64 = left.astype(np.float64), right.astype(np.float64) + denominator = np.maximum(np.abs(left64), 1e-8) + max_rel = max(max_rel, float(np.max(np.abs(left64-right64)/denominator))) + checked += left.size + return {"max_abs_error":max_abs, "max_rel_error":max_rel, "checked_values":checked} + + +def validate(candidate_path: Path, config: HarnessConfig, prompt_lengths: tuple[int, ...], parameters: Mapping[str, Any]) -> dict[str, Any]: + block = _block(config) + model = ModelShell(block, config) + factory = _load_candidate(candidate_path, parameters) + max_abs = max_rel = 0.0 + checked = 0 + cases = [] + for case_index, prompt_tokens in enumerate(prompt_lengths): + _reset_cache(block) + reference = _run_transition(block, config, prompt_tokens, seed=1000+case_index) + _reset_cache(block) + context = {"stage":"transition_oracle", "prompt_tokens":prompt_tokens, "batch_size":1, + "parameters":dict(parameters), "device":Device.DEFAULT} + replacement = factory(block, 0, model, context) + if not callable(replacement): raise TypeError("build_replacement must return a callable block") + candidate = _run_transition(replacement, config, prompt_tokens, seed=1000+case_index) + result = _error(reference, candidate) + max_abs, max_rel = max(max_abs, result["max_abs_error"]), max(max_rel, result["max_rel_error"]) + checked += result["checked_values"] + cases.append({"prompt_tokens":prompt_tokens, **result}) + passed = max_abs <= config.max_abs_error and max_rel <= config.max_rel_error + return {"passed":passed, "max_abs_error":max_abs, "max_rel_error":max_rel, "checked_values":checked, + "reason":"" if passed else "candidate diverged from transformer-block transition reference", "cases":cases} + + +def benchmark(candidate_path: Path, config: HarnessConfig, budget: int, parameters: Mapping[str, Any]) -> tuple[list[float], dict[str, Any]]: + block = _block(config) + model = ModelShell(block, config) + factory = _load_candidate(candidate_path, parameters) + replacement = factory(block, 0, model, {"stage":"decode_benchmark", "parameters":dict(parameters), "device":Device.DEFAULT}) + if not callable(replacement): raise TypeError("build_replacement must return a callable block") + _reset_cache(block) + prompt, _, _, freqs, mask = _inputs(config, config.prompt_tokens, seed=2026) + replacement(prompt, 0, freqs[:, :config.prompt_tokens], mask).realize() + samples = [] + kernel_counts = [] + for index in range(budget): + position = config.prompt_tokens + index + Tensor.manual_seed(3000+index) + value = Tensor.randn(1, 1, config.dim).cast(_dtype(config.dtype)).contiguous().realize() + GlobalCounters.reset() + started = time.perf_counter_ns() + replacement(value, position, freqs[:, position:position+1], None).realize() + samples.append((time.perf_counter_ns()-started)/1000.0) + kernel_counts.append(GlobalCounters.kernel_count) + return samples, {"kernel_count_mean":sum(kernel_counts)/len(kernel_counts), "device":Device.DEFAULT, + "prompt_tokens":config.prompt_tokens, "stage":"decode", "parameters":dict(parameters)} + + +def main() -> None: + parser = argparse.ArgumentParser(description="Radeon Forge transformer-block oracle and benchmark") + parser.add_argument("--candidate") + parser.add_argument("--mode", choices=("correctness", "benchmark", "transition"), default="correctness") + parser.add_argument("--dtype", choices=("float32", "float16", "bfloat16"), default="float32") + parser.add_argument("--dim", type=int, default=128) + parser.add_argument("--hidden-dim", type=int, default=256) + parser.add_argument("--prompt-tokens", type=int, default=8) + args = parser.parse_args() + config = HarnessConfig(dim=args.dim, hidden_dim=args.hidden_dim, prompt_tokens=args.prompt_tokens, dtype=args.dtype) + candidate_path, parameters = _candidate_path(args.candidate), _parameters() + prompt_lengths = (1, 4, args.prompt_tokens, min(24, config.max_context-2)) if args.mode == "transition" else (args.prompt_tokens,) + correctness = validate(candidate_path, config, tuple(dict.fromkeys(prompt_lengths)), parameters) + budget = max(1, int(os.environ.get("RADEON_FORGE_BUDGET", "5"))) + samples, metrics = (benchmark(candidate_path, config, budget, parameters) if args.mode == "benchmark" and correctness["passed"] else ([], {})) + payload = { + "samples_us":samples, + "correctness":correctness, + "resources":{}, + "compile_ok":True, + "stable":bool(correctness["passed"]), + "metrics":metrics, + "mode":args.mode, + "candidate_sha256":hashlib.sha256(candidate_path.read_bytes()).hexdigest(), + } + print(json.dumps(payload, default=str)) + raise SystemExit(0 if correctness["passed"] else 2) + + +if __name__ == "__main__": main() diff --git a/extra/radeon_forge/workloads/private_code_agent_suite.json b/extra/radeon_forge/workloads/private_code_agent_suite.json new file mode 100644 index 0000000000000..2c14f3438bb0a --- /dev/null +++ b/extra/radeon_forge/workloads/private_code_agent_suite.json @@ -0,0 +1,58 @@ +{ + "format_version": 1, + "name": "private-code-agent-v1", + "description": "Deterministic local repository tasks covering inspection, retrieval, tool execution, and resumed generation.", + "metadata": { + "primary_metric": "p95_end_to_end_task_latency", + "constraints": { + "external_network_calls": 0, + "tool_call_validity": 1.0, + "task_success_drop_from_baseline": 0.0 + } + }, + "tasks": [ + { + "task_id": "read-project-identity", + "prompt": "Read the first 80 lines of README.md and tell me the project name. Use the local file tool before answering.", + "expected_tools": ["read_file"], + "allowed_tools": ["read_file"], + "required_output_regex": ["tinygrad"], + "forbidden_output_regex": ["https?://"], + "max_tokens": 192, + "timeout_seconds": 120 + }, + { + "task_id": "locate-global-counters", + "prompt": "Find where GlobalCounters is defined in this private repository, then answer with the file path. Use local text search.", + "expected_tools": ["search_text"], + "allowed_tools": ["search_text"], + "required_output_regex": ["tinygrad/helpers\\.py"], + "forbidden_output_regex": ["https?://"], + "max_tokens": 192, + "timeout_seconds": 120 + }, + { + "task_id": "run-forge-unit-test", + "prompt": "Run only the Radeon Forge KV ledger unit test and report whether it passed. Do not run the full suite.", + "expected_tools": ["run_command"], + "allowed_tools": ["run_command"], + "required_output_regex": ["pass|ok"], + "forbidden_output_regex": ["https?://"], + "max_tokens": 192, + "timeout_seconds": 180, + "metadata": { + "expected_command": "python -m unittest test.test_radeon_forge_kv" + } + }, + { + "task_id": "explain-local-profile-boundary", + "prompt": "Without calling a tool, explain in one sentence why MockGPU timing cannot be used as Radeon performance evidence.", + "expected_tools": [], + "allowed_tools": [], + "required_output_regex": ["hardware|W7900|real GPU"], + "forbidden_output_regex": ["https?://"], + "max_tokens": 96, + "timeout_seconds": 90 + } + ] +} diff --git a/extra/radeon_forge/workloads/rdna3_asm_matmul.py b/extra/radeon_forge/workloads/rdna3_asm_matmul.py new file mode 100644 index 0000000000000..2079e112d77d7 --- /dev/null +++ b/extra/radeon_forge/workloads/rdna3_asm_matmul.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import json +import math +import os +import statistics + +import numpy as np + +from tinygrad import Context, Device, GlobalCounters, Tensor +from tinygrad.dtype import AddrSpace, dtypes +from tinygrad.engine.realize import Estimates, run_linear +from tinygrad.uop.ops import KernelInfo, Ops, UOp + +from extra.gemm import amd_asm_matmul as gemm + + +_ORDER_BUILDERS = { + "optimized": lambda: list(gemm.FMAC_PAIR_ORDER), + "row_major": lambda: [(a, b) for a in range(4) for b in range(8)], + "column_major": lambda: [(a, b) for b in range(8) for a in range(4)], + "snake": lambda: [(a, b) for a in range(4) for b in (range(8) if a % 2 == 0 else range(7, -1, -1))], +} + + +def _candidate() -> tuple[dict, int]: + candidate = json.loads(os.environ["RADEON_FORGE_CANDIDATE_JSON"]) + budget = int(os.environ.get("RADEON_FORGE_BUDGET", "5")) + if budget <= 0: raise ValueError("RADEON_FORGE_BUDGET must be positive") + return candidate, budget + + +def _install_order(name: str) -> None: + if name not in _ORDER_BUILDERS: raise ValueError(f"unknown FMAC order {name!r}") + order = _ORDER_BUILDERS[name]() + expected = {(a, b) for a in range(4) for b in range(8)} + if len(order) != 32 or set(order) != expected: raise ValueError(f"invalid FMAC order {name!r}") + gemm.FMAC_PAIR_ORDER = order + gemm.FMAC_PATTERN = gemm.derive_fmac_pattern(gemm.ACC_GRID, gemm.V_A_TILE_REGS, gemm.V_B_TILE_REGS) + + +def _percentile(values: list[float], q: float) -> float: + ordered = sorted(values) + return ordered[min(len(ordered) - 1, max(0, math.ceil(q * len(ordered)) - 1))] + + +def run() -> dict: + candidate, budget = _candidate() + parameters = candidate["parameters"] + n = int(parameters["N"]) + limit_occ = int(parameters["LIMIT_OCC"]) + order_name = str(parameters["FMAC_ORDER"]) + if limit_occ <= 0: raise ValueError("LIMIT_OCC must be positive") + _install_order(order_name) + + dev = Device[Device.DEFAULT] + arch = getattr(getattr(dev, "renderer", None), "target", None) + arch_name = getattr(arch, "arch", None) + if arch_name != "gfx1100": raise RuntimeError(f"RDNA3 workload requires gfx1100, got {arch_name!r}") + + instructions = gemm.build_kernel(n) + rng = np.random.default_rng(42) + a = Tensor(rng.random((n, n), dtype=np.float32) - 0.5) + b = Tensor(rng.random((n, n), dtype=np.float32) - 0.5) + c = Tensor.empty(n, n) + Tensor.realize(a, b, c) + + grid, local = (n // gemm.BLOCK_N, n // gemm.BLOCK_M, 1), (gemm.THREADS, 1, 1) + lds_size = max(gemm.LDS_SIZE, 65536 // limit_occ) + + def asm_kernel(A: UOp, B: UOp, C: UOp) -> UOp: + gidxs = [UOp.special(size, f"gidx{index}") for index, size in enumerate(grid)] + lidxs = [UOp.special(size, f"lidx{index}") for index, size in enumerate(local)] + lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL) + sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs, + arg=KernelInfo(name=f"radeon_forge_{candidate['candidate_id']}", estimates=Estimates(ops=n*n*n*2, mem=n*n*4*3))) + return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=instruction) for instruction in instructions)))) + + c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2] + linear = c.schedule_linear() + with Context(DEBUG=0): + run_linear(linear) # warmup and realize all lazy dependencies + samples_us: list[float] = [] + for _ in range(budget): + start = GlobalCounters.time_sum_s + run_linear(linear) + elapsed = (GlobalCounters.time_sum_s - start) * 1e6 + if not math.isfinite(elapsed) or elapsed <= 0: raise RuntimeError(f"invalid timing sample {elapsed}") + samples_us.append(elapsed) + + reference = (a @ b).realize() + output_np, reference_np = c.numpy(), reference.numpy() + difference = np.abs(output_np - reference_np) + denominator = np.maximum(np.abs(reference_np), 1e-6) + max_abs_error = float(difference.max()) + max_rel_error = float((difference / denominator).max()) + mse = float(np.square(output_np - reference_np).mean()) + stable = len(samples_us) < 3 or statistics.pstdev(samples_us) / statistics.mean(samples_us) <= 0.10 + + return { + "compile_ok": True, + "stable": stable, + "samples_us": samples_us, + "median_latency_us": statistics.median(samples_us), + "p95_latency_us": _percentile(samples_us, 0.95), + "correctness": { + "passed": bool(np.isfinite(mse) and mse <= 1e-6), + "max_abs_error": max_abs_error, + "max_rel_error": max_rel_error, + "checked_values": int(output_np.size), + "reason": "" if np.isfinite(mse) and mse <= 1e-6 else f"mean squared error {mse}", + }, + "resources": { + "vgprs": 179, + "sgprs": 56, + "lds_bytes": lds_size, + "scratch_bytes": 0, + "spilled_vgprs": 0, + "spilled_sgprs": 0, + }, + "evidence": { + "target": arch_name, + "reference": "tinygrad matmul", + "mean_squared_error": mse, + "fmac_order": order_name, + "limit_occ": limit_occ, + "resource_source": "static register assignments and explicit LDS allocation", + "kernel_instruction_count": len(instructions), + }, + } + + +if __name__ == "__main__": print(json.dumps(run(), sort_keys=True)) diff --git a/test/test_radeon_forge.py b/test/test_radeon_forge.py new file mode 100644 index 0000000000000..5135d95352abd --- /dev/null +++ b/test/test_radeon_forge.py @@ -0,0 +1,140 @@ +import json +import sys +import tempfile +import unittest +from pathlib import Path + +from extra.radeon_forge.agent import AgentState, AgentStateError, ForgeAgent +from extra.radeon_forge.amd_metadata import resource_report_from_comgr, resource_report_from_text +from extra.radeon_forge.command_backend import CommandHarness +from extra.radeon_forge.contracts import Candidate, CorrectnessReport, ResourceReport, TrialResult, WorkloadContract +from extra.radeon_forge.families import rdna3_asm_matmul_family, rdna3_rmsnorm_fp8_family +from extra.radeon_forge.knowledge import LocalKnowledgeBase +from extra.radeon_forge.ledger import ExperimentLedger +from extra.radeon_forge.permissions import Action, PermissionController, PermissionDenied +from extra.radeon_forge.planner import LocalEndpointRequired, PlannerDecision, validate_local_endpoint +from extra.radeon_forge.tuner import successive_halving + + +class FakePlanner: + def plan(self, request, contract, evidence=(), memory=()): + return PlannerDecision( + summary=f"Tune {contract.name}", + hypothesis="A different FMAC order may reduce exposed stalls.", + candidate_family="rdna3-asm-matmul", + proposed_actions=(Action.BENCHMARK,), + benchmark_budget=2, + rationale=f"Measure on {contract.target}; evidence={len(evidence)}, memory={len(memory)}", + ) + + +class TestRadeonForge(unittest.TestCase): + def test_metadata_text(self): + text = """ +; NumSgprs: 18 +; NumVgprs: 54 +; ScratchSize: 0 +; LDSByteSize: 8192 bytes/workgroup +; Occupancy: 8 + .vgpr_spill_count: 0 + .sgpr_spill_count: 1 +""" + report = resource_report_from_text(text) + self.assertEqual((report.sgprs, report.vgprs, report.lds_bytes, report.occupancy), (18, 54, 8192, 8)) + self.assertTrue(report.has_spills) + + def test_metadata_comgr(self): + report = resource_report_from_comgr({"amdhsa.kernels": [{".vgpr_count": "44", ".sgpr_count": "20", + ".group_segment_fixed_size": "4096", ".private_segment_fixed_size": "0", ".vgpr_spill_count": "0", ".sgpr_spill_count": "0"}]}) + self.assertEqual((report.vgprs, report.sgprs, report.lds_bytes), (44, 20, 4096)) + self.assertFalse(report.has_spills) + + def test_family_grids(self): + rmsnorm = rdna3_rmsnorm_fp8_family("/tmp/tinygrad", 4096 * 16, 4096).candidates() + self.assertEqual(len(rmsnorm), 12) + self.assertEqual(len({candidate.candidate_id for candidate in rmsnorm}), 12) + self.assertTrue(all(candidate.parameters["HIDDEN"] == 4096 for candidate in rmsnorm)) + matmul = rdna3_asm_matmul_family("/tmp/tinygrad", 1024).candidates() + self.assertEqual(len(matmul), 12) + self.assertEqual({candidate.parameters["FMAC_ORDER"] for candidate in matmul}, {"optimized", "row_major", "column_major", "snake"}) + + def test_hard_gate_beats_fast_invalid_candidate(self): + candidates = [Candidate("bad", "test", {}), Candidate("good", "test", {})] + contract = WorkloadContract("unit", max_vgprs=96) + + def evaluate(candidate, budget): + if candidate.candidate_id == "bad": + return TrialResult(candidate, CorrectnessReport(True), samples_us=(1.0,) * budget, median_latency_us=1.0, p95_latency_us=1.0, + resources=ResourceReport(vgprs=128)) + return TrialResult(candidate, CorrectnessReport(True), samples_us=(2.0,) * budget, median_latency_us=2.0, p95_latency_us=2.0, + resources=ResourceReport(vgprs=64)) + + with tempfile.TemporaryDirectory() as directory: + ledger = ExperimentLedger(Path(directory) / "events.jsonl") + summary = successive_halving(candidates, evaluate, contract, budgets=(1, 2), reduction=2, ledger=ledger) + self.assertIsNotNone(summary.winner) + self.assertEqual(summary.winner.candidate.candidate_id, "good") + self.assertTrue(any(record["event"] == "winner_selected" for record in ledger.records())) + + def test_permissions_are_scoped_and_single_use(self): + controller = PermissionController() + with self.assertRaises(PermissionDenied): controller.authorize(None, Action.BENCHMARK) + grant = controller.issue([Action.BENCHMARK], "run two short hardware trials", max_uses=1) + self.assertEqual(controller.authorize(grant.token, Action.BENCHMARK).reason, grant.reason) + with self.assertRaises(PermissionDenied): controller.authorize(grant.token, Action.BENCHMARK) + deploy_grant = controller.issue([Action.BENCHMARK], "benchmark only") + with self.assertRaises(PermissionDenied): controller.authorize(deploy_grant.token, Action.DEPLOY) + + def test_command_harness_json_protocol(self): + candidate = Candidate("candidate", "unit", {"THREADS": 128}) + with tempfile.TemporaryDirectory() as directory: + runner = Path(directory) / "runner.py" + payload = { + "samples_us": [3.0, 2.0, 4.0], + "correctness": {"passed": True, "checked_values": 16}, + "resources": {"vgprs": 32, "spilled_vgprs": 0, "spilled_sgprs": 0}, + } + runner.write_text(f"import json\nprint(json.dumps({json.dumps(payload)}))\n", encoding="utf-8") + result = CommandHarness([sys.executable, str(runner)])(candidate, 3) + self.assertTrue(result.correctness.passed) + self.assertEqual(result.median_latency_us, 3.0) + self.assertEqual(result.p95_latency_us, 4.0) + self.assertEqual(result.resources.vgprs, 32) + + def test_planner_rejects_remote_endpoints(self): + self.assertEqual(validate_local_endpoint("http://127.0.0.1:8000/v1"), "http://127.0.0.1:8000/v1") + self.assertEqual(validate_local_endpoint("http://[::1]:8000/v1"), "http://[::1]:8000/v1") + with self.assertRaises(LocalEndpointRequired): validate_local_endpoint("https://api.example.com/v1") + with self.assertRaises(LocalEndpointRequired): validate_local_endpoint("http://192.168.1.5:8000/v1") + + def test_local_knowledge_retrieval_has_citations(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "knowledge.md" + path.write_text("RDNA3 uses VOPD scheduling.\nSpilled VGPRs are rejected.\nRMSNorm can remove HBM materialization.\n", encoding="utf-8") + knowledge = LocalKnowledgeBase.from_paths([path], lines_per_chunk=2, overlap=1) + hits = knowledge.search("VGPR spills", top_k=2) + self.assertTrue(hits) + self.assertIn(str(path), hits[0].citation) + self.assertIn("VGPR", hits[0].chunk.text) + + def test_multi_turn_agent_requires_matching_approval(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "knowledge.md" + path.write_text("RDNA3 VOPD ordering is benchmark-sensitive.\n", encoding="utf-8") + knowledge = LocalKnowledgeBase.from_paths([path]) + ledger = ExperimentLedger(Path(directory) / "events.jsonl") + agent = ForgeAgent(WorkloadContract("private-code-agent"), FakePlanner(), knowledge, ledger) + proposal = agent.propose("Optimize P95 latency without changing precision") + self.assertEqual(agent.state, AgentState.AWAITING_APPROVAL) + with self.assertRaises(AgentStateError): agent.approve("wrong-id", "approved") + agent.approve(proposal.proposal_id, "I approve two local benchmark trials", max_tool_uses=2) + agent.authorize(Action.BENCHMARK) + agent.complete({"status": "accepted", "speedup": 1.1}) + self.assertEqual(agent.state, AgentState.COMPLETED) + events = [record["event"] for record in ledger.records()] + self.assertIn("proposal_created", events) + self.assertIn("proposal_approved", events) + self.assertIn("execution_completed", events) + + +if __name__ == "__main__": unittest.main() diff --git a/test/test_radeon_forge_autotune.py b/test/test_radeon_forge_autotune.py new file mode 100644 index 0000000000000..06b021084820b --- /dev/null +++ b/test/test_radeon_forge_autotune.py @@ -0,0 +1,54 @@ +import json +import sys +import tempfile +import unittest +from pathlib import Path + +from extra.radeon_forge.synthesis.autotune import SearchPlan, run_autotune +from extra.radeon_forge.synthesis.workspace import CandidateWorkspace, KernelSpec + + +class TestStageAutotune(unittest.TestCase): + def test_fast_invalid_configuration_cannot_win(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + runner = root / "runner.py" + runner.write_text('''import json, os +candidate = json.loads(os.environ["RADEON_FORGE_CANDIDATE_JSON"]) +tile = candidate["parameters"]["TILE"] +budget = int(os.environ["RADEON_FORGE_BUDGET"]) +payload = { + "samples_us": ([1.0] if tile == 1 else [2.0]) * budget, + "correctness": {"passed": True, "max_abs_error": 0.0, "max_rel_error": 0.0, "checked_values": 128}, + "resources": {"vgprs": 128 if tile == 1 else 64, "lds_bytes": 4096, "spilled_vgprs": 0, "spilled_sgprs": 0}, + "compile_ok": True, + "stable": True +} +print(json.dumps(payload)) +''', encoding="utf-8") + workspace = CandidateWorkspace(root / "forge") + spec = workspace.save_spec(KernelSpec( + name="decode-search", operation="batch-one decode subgraph", target="gfx1100", invariants=("match reference",), + hardware_command=(sys.executable, str(runner)), + metadata={"recipe_acceptance":{"max_vgprs":96, "forbid_spills":True}, + "hook":{"layer":"subgraph", "target":"decode.projection", "adapter":"request_metadata", + "when":{"stages":["decode"], "batch_sizes":[1]}, "exclusive_group":"decode-projection"}}, + )) + candidate = workspace.create_candidate(spec.spec_id, "# parameterized candidate\n", "search tile size") + workspace.update(candidate.candidate_id, "mockgpu_passed", {"mockgpu":{"passed":True}}) + updated, summary = run_autotune(workspace, candidate.candidate_id, root, + SearchPlan.from_mapping({"axes":{"TILE":[1,2]}, "budgets":[1,2], "reduction":2})) + self.assertIsNotNone(summary.winner) + self.assertEqual(summary.winner.candidate.parameters["TILE"], 2) + self.assertEqual(updated.status, "hardware_passed") + self.assertEqual(updated.evidence["selected_parameters"], {"TILE":2}) + persisted = workspace.load_candidate(candidate.candidate_id) + self.assertEqual(persisted.evidence["selected_parameters"], {"TILE":2}) + self.assertTrue(Path(persisted.evidence["autotune"]["ledger"]).is_file()) + + def test_search_space_is_bounded(self): + with self.assertRaisesRegex(ValueError, "maximum"): + SearchPlan.from_mapping({"axes":{"A":list(range(20)), "B":list(range(20))}, "max_candidates":128}) + + +if __name__ == "__main__": unittest.main() diff --git a/test/test_radeon_forge_block_harness.py b/test/test_radeon_forge_block_harness.py new file mode 100644 index 0000000000000..ae4ac277bacf4 --- /dev/null +++ b/test/test_radeon_forge_block_harness.py @@ -0,0 +1,55 @@ +import tempfile +import unittest +from pathlib import Path + +from extra.radeon_forge.workloads.block_hook_harness import HarnessConfig, benchmark, validate + + +IDENTITY = ''' +PARAMETERS = {} +def configure(parameters): + global PARAMETERS + PARAMETERS = dict(parameters) +def build_replacement(original, layer_index, model, context): + assert context["parameters"] == PARAMETERS + return original +''' + +BAD = ''' +def build_replacement(original, layer_index, model, context): + def replacement(x, start_pos, freqs_cis, mask): + return x * 0 + return replacement +''' + + +class TestBlockHookHarness(unittest.TestCase): + def test_identity_replacement_passes_transition_and_benchmark(self): + with tempfile.TemporaryDirectory() as directory: + path=Path(directory)/"identity.py" + path.write_text(IDENTITY,encoding="utf-8") + config=HarnessConfig(dim=32,hidden_dim=64,n_heads=4,n_kv_heads=4,max_context=32,prompt_tokens=2, + max_abs_error=1e-5,max_rel_error=1e-4) + result=validate(path,config,(1,2,4),{"WAVES":2}) + self.assertTrue(result["passed"],result) + self.assertEqual(len(result["cases"]),3) + self.assertGreater(result["checked_values"],0) + samples,metrics=benchmark(path,config,2,{"WAVES":2}) + self.assertEqual(len(samples),2) + self.assertTrue(all(value>0 for value in samples)) + self.assertEqual(metrics["stage"],"decode") + self.assertEqual(metrics["parameters"],{"WAVES":2}) + + def test_incorrect_replacement_fails_reference_oracle(self): + with tempfile.TemporaryDirectory() as directory: + path=Path(directory)/"bad.py" + path.write_text(BAD,encoding="utf-8") + config=HarnessConfig(dim=32,hidden_dim=64,n_heads=4,n_kv_heads=4,max_context=16,prompt_tokens=2, + max_abs_error=1e-5,max_rel_error=1e-4) + result=validate(path,config,(2,),{}) + self.assertFalse(result["passed"]) + self.assertGreater(result["max_abs_error"],config.max_abs_error) + self.assertIn("diverged",result["reason"]) + + +if __name__=="__main__": unittest.main() diff --git a/test/test_radeon_forge_capture.py b/test/test_radeon_forge_capture.py new file mode 100644 index 0000000000000..eeb55ddcfba2b --- /dev/null +++ b/test/test_radeon_forge_capture.py @@ -0,0 +1,89 @@ +import os +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +from extra.radeon_forge.profiling.capture import CaptureKind, CaptureRequest, Rocprofv3Adapter +from extra.radeon_forge.profiling.tools import ProfilingTools +from extra.radeon_forge.runtime.tool_context import ToolExecutionContext, bind_tool_context + + +FAKE_ROCPROF = '''#!/usr/bin/env python3 +import csv, pathlib, sys +if "--version" in sys.argv: + print("rocprofv3 fake 1.0") + raise SystemExit(0) +if "--help" in sys.argv: + print("--att --att-simd-select --kernel-include-regex --pmc --output-directory -d") + raise SystemExit(0) +args=sys.argv[1:] +flag="--output-directory" if "--output-directory" in args else "-d" +out=pathlib.Path(args[args.index(flag)+1]) +out.mkdir(parents=True, exist_ok=True) +if "--att" in args: + (out/"thread_trace.att").write_bytes(b"ATT-EVIDENCE") +else: + with (out/"counters.csv").open("w", newline="") as f: + writer=csv.DictWriter(f, fieldnames=["Kernel_Name","SQ_WAVES","DurationNs"]) + writer.writeheader() + writer.writerow({"Kernel_Name":"decode_gemv","SQ_WAVES":"64","DurationNs":"900"}) +print("capture complete") +''' + + +class TestRocprofCapture(unittest.TestCase): + def _adapter(self, root: Path): + binary=root/"bin"/"rocprofv3" + binary.parent.mkdir() + binary.write_text(FAKE_ROCPROF, encoding="utf-8") + binary.chmod(0o755) + environment={**os.environ, "PATH":str(binary.parent)+os.pathsep+os.environ.get("PATH","")} + return Rocprofv3Adapter(root, root/"evidence"), environment + + def test_counter_capture_writes_manifest_hashes_and_summary(self): + with tempfile.TemporaryDirectory() as directory: + root=Path(directory) + adapter, environment=self._adapter(root) + with patch.dict(os.environ, environment, clear=True): + probe=adapter.probe() + self.assertTrue(probe["available"]) + self.assertTrue(probe["supports_pmc"]) + result=adapter.capture(CaptureRequest(CaptureKind.COUNTERS, ("python3","-c","print('x')"), "decode", + session_id="s1", trace_id="t1", agent_step=2, counters=("SQ_WAVES","DurationNs"))) + self.assertTrue(result.passed) + self.assertEqual(result.summary["stage"], "decode") + self.assertEqual(result.summary["session_id"], "s1") + self.assertEqual(result.summary["parsed_csv"]["csv_rows"], 1) + self.assertEqual(result.summary["parsed_csv"]["numeric_columns"]["SQ_WAVES"]["mean"], 64.0) + self.assertTrue(all(len(artifact.sha256)==64 for artifact in result.artifacts)) + self.assertTrue((Path(result.output_directory)/"forge_capture_manifest.json").is_file()) + + def test_att_capture_is_separate_and_provenance_is_bound(self): + with tempfile.TemporaryDirectory() as directory: + root=Path(directory) + adapter, environment=self._adapter(root) + tools=ProfilingTools(root, root/"evidence") + context=ToolExecutionContext("session-7","trace-9",3,"call-2","capture_rocm_att",{"purpose":"decode diagnosis"}) + with patch.dict(os.environ, environment, clear=True), bind_tool_context(context): + result=tools.capture_att({"command":["python3","-c","print('decode')"], "stage":"first_token", + "kernel_regex":"decode_.*", "timeout_seconds":30}) + self.assertTrue(result["passed"]) + self.assertEqual(result["summary"]["session_id"], "session-7") + self.assertEqual(result["summary"]["trace_id"], "trace-9") + self.assertEqual(result["summary"]["agent_step"], 3) + self.assertEqual(result["summary"]["kernel_regex"], "decode_.*") + self.assertTrue(any(item["suffix"]==".att" for item in result["artifacts"])) + listed=tools.list_captures({"limit":10}) + self.assertEqual(listed["captures"][0]["capture_id"], result["capture_id"]) + + def test_profile_target_is_allowlisted_and_shell_free(self): + with tempfile.TemporaryDirectory() as directory: + root=Path(directory) + adapter, environment=self._adapter(root) + with patch.dict(os.environ, environment, clear=True): + with self.assertRaisesRegex(ValueError, "not allowlisted"): + adapter.capture(CaptureRequest(CaptureKind.ATT, ("bash","-lc","rm -rf /"), "decode")) + + +if __name__ == "__main__": unittest.main() diff --git a/test/test_radeon_forge_deployment_metrics.py b/test/test_radeon_forge_deployment_metrics.py new file mode 100644 index 0000000000000..93921e2963e41 --- /dev/null +++ b/test/test_radeon_forge_deployment_metrics.py @@ -0,0 +1,52 @@ +import tempfile +import unittest + +from extra.radeon_forge.synthesis.deployment import SafeHookRegistry +from extra.radeon_forge.synthesis.hooks import HookActivationError,RuntimeFingerprint +from extra.radeon_forge.synthesis.workspace import CandidateWorkspace,KernelSpec + + +SOURCE=''' +def build_replacement(original, layer_index, model, context): + return original +''' + + +class TestDeploymentMetricGate(unittest.TestCase): + def test_stateful_candidate_requires_transition_and_agent_metric(self): + with tempfile.TemporaryDirectory() as directory: + workspace=CandidateWorkspace(directory) + spec=workspace.save_spec(KernelSpec("decode","block",target="gfx1100",invariants=("match",), + metadata={"heldout_command":["python3","transition.py"],"require_agent_evaluation":True, + "hook":{"layer":"transformer_block","target":"llama.block","adapter":"python_transformer_block", + "when":{"stages":["first_token","decode"]}}})) + candidate=workspace.create_candidate(spec.spec_id,SOURCE,"faster block") + workspace.update(candidate.candidate_id,"heldout_passed",{"heldout":{"passed":True}}) + registry=SafeHookRegistry(workspace) + fingerprint=RuntimeFingerprint(architecture="gfx1100",runtime="tinygrad",model_family="llama") + with self.assertRaisesRegex(HookActivationError,"frozen agent-suite"): + registry.activate(candidate.candidate_id,fingerprint,"deploy") + + workspace.update(candidate.candidate_id,"evaluation_failed",{"agent_evaluation_comparison":{"gate":{"passed":False}}}) + with self.assertRaisesRegex(HookActivationError,"frozen agent-suite"): + registry.activate(candidate.candidate_id,fingerprint,"deploy") + + workspace.update(candidate.candidate_id,"evaluation_passed",{"agent_evaluation_comparison":{"gate":{"passed":True}}}) + active=registry.activate(candidate.candidate_id,fingerprint,"quality and latency gates passed") + self.assertEqual(active.candidate_id,candidate.candidate_id) + + def test_stateless_non_final_hook_keeps_lighter_policy(self): + with tempfile.TemporaryDirectory() as directory: + workspace=CandidateWorkspace(directory) + spec=workspace.save_spec(KernelSpec("advisory","block",target="gfx1100",invariants=("match",), + metadata={"heldout_command":["python3","transition.py"], + "hook":{"layer":"transformer_block","target":"llama.block","adapter":"python_transformer_block", + "when":{"stages":["decode"]}}})) + candidate=workspace.create_candidate(spec.spec_id,SOURCE,"validated") + workspace.update(candidate.candidate_id,"heldout_passed",{}) + active=SafeHookRegistry(workspace).activate(candidate.candidate_id, + RuntimeFingerprint(architecture="gfx1100"),"transition passed") + self.assertEqual(active.candidate_id,candidate.candidate_id) + + +if __name__=="__main__":unittest.main() diff --git a/test/test_radeon_forge_evaluation.py b/test/test_radeon_forge_evaluation.py new file mode 100644 index 0000000000000..f3af4c1b13b02 --- /dev/null +++ b/test/test_radeon_forge_evaluation.py @@ -0,0 +1,59 @@ +import json +import tempfile +import unittest +from pathlib import Path + +from extra.radeon_forge.backends.fake import ScriptedBackend +from extra.radeon_forge.evaluation.suite import AgentTask, AgentTaskSuite, run_suite +from extra.radeon_forge.runtime import ForgeEngine + + +class TestAgentEvaluation(unittest.TestCase): + def test_suite_measures_complete_tool_loop(self): + responses=[ + '{"name":"read_file","arguments":{"path":"marker.txt"}}', + 'The marker is ALPHA.', + 'MockGPU is a semantic emulator; only real W7900 hardware provides performance evidence.', + ] + with tempfile.TemporaryDirectory() as directory: + Path(directory,"marker.txt").write_text("ALPHA",encoding="utf-8") + suite=AgentTaskSuite("unit",( + AgentTask("read","Read marker.txt and report the marker.",expected_tools=("read_file",),allowed_tools=("read_file",), + required_output_regex=("ALPHA",),timeout_seconds=10), + AgentTask("explain","Why not use MockGPU timing?",required_output_regex=("real W7900|real.*hardware",),timeout_seconds=10), + )) + engine=ForgeEngine(ScriptedBackend(responses),directory) + result=run_suite(engine,suite) + self.assertEqual(result.passed_tasks,2) + self.assertEqual(result.task_success_rate,1.0) + self.assertEqual(result.tool_call_validity_rate,1.0) + self.assertGreater(result.p50_latency_ms,0) + self.assertEqual(result.results[0].observed_tools,("read_file",)) + self.assertTrue(result.results[0].trace_id) + engine.close() + + def test_unexpected_tool_is_rejected_and_task_fails(self): + with tempfile.TemporaryDirectory() as directory: + suite=AgentTaskSuite("unit",(AgentTask("bad","Answer directly",expected_tools=(),allowed_tools=()),)) + backend=ScriptedBackend(['{"name":"list_files","arguments":{}}']) + engine=ForgeEngine(backend,directory) + result=run_suite(engine,suite) + task=result.results[0] + self.assertFalse(task.passed) + self.assertFalse(task.tool_call_valid) + self.assertIn("not allowed",task.failure_reasons[0]) + self.assertEqual(task.terminal_state,"idle") + engine.close() + + def test_suite_hash_is_stable_and_file_load_validates_ids(self): + suite=AgentTaskSuite("stable",(AgentTask("a","prompt"),),metadata={"x":1}) + self.assertEqual(suite.suite_hash,suite.suite_hash) + with tempfile.TemporaryDirectory() as directory: + path=Path(directory)/"suite.json" + payload={"format_version":1,"name":"x","tasks":[{"task_id":"a","prompt":"one"},{"task_id":"a","prompt":"two"}]} + path.write_text(json.dumps(payload),encoding="utf-8") + with self.assertRaisesRegex(ValueError,"unique"): + AgentTaskSuite.load(path) + + +if __name__=="__main__":unittest.main() diff --git a/test/test_radeon_forge_evaluation_compare.py b/test/test_radeon_forge_evaluation_compare.py new file mode 100644 index 0000000000000..7ca4714f4a6c2 --- /dev/null +++ b/test/test_radeon_forge_evaluation_compare.py @@ -0,0 +1,46 @@ +import unittest + +from extra.radeon_forge.evaluation.compare import IncomparableSuites,compare_evaluations + + +def result(suite_hash="same",success=1.0,tools=1.0,p95=100.0,passed=(True,True),latencies=(80.0,120.0)): + return {"suite_name":"private","suite_hash":suite_hash,"task_success_rate":success,"tool_call_validity_rate":tools, + "p50_latency_ms":90.0,"p95_latency_ms":p95,"mean_latency_ms":100.0, + "results":[ + {"task_id":"a","passed":passed[0],"latency_ms":latencies[0],"observed_tools":["read_file"],"failure_reasons":[]}, + {"task_id":"b","passed":passed[1],"latency_ms":latencies[1],"observed_tools":[],"failure_reasons":[]}, + ]} + + +class TestEvaluationComparison(unittest.TestCase): + def test_faster_quality_preserving_candidate_is_deployable(self): + baseline=result(p95=120.0) + candidate=result(p95=80.0,latencies=(60.0,90.0)) + comparison=compare_evaluations(baseline,candidate) + self.assertTrue(comparison["gate"]["passed"]) + self.assertTrue(comparison["deployable_speedup"]) + self.assertEqual(comparison["latency"]["p95_latency_ms"]["absolute"],-40.0) + self.assertEqual(len(comparison["per_task"]),2) + + def test_speed_cannot_hide_task_or_tool_regression(self): + baseline=result(p95=120.0) + candidate=result(success=0.5,tools=0.5,p95=50.0,passed=(True,False),latencies=(30.0,40.0)) + comparison=compare_evaluations(baseline,candidate) + self.assertFalse(comparison["gate"]["passed"]) + self.assertFalse(comparison["deployable_speedup"]) + self.assertIn("task success regressed",comparison["gate"]["reasons"][0]) + self.assertTrue(any("baseline-passing tasks regressed" in reason for reason in comparison["gate"]["reasons"])) + + def test_different_suite_or_task_ids_are_rejected(self): + with self.assertRaisesRegex(IncomparableSuites,"suite_hash differs"): + compare_evaluations(result("a"),result("b")) + candidate=result();candidate["results"]=candidate["results"][:1] + with self.assertRaisesRegex(IncomparableSuites,"task ids differ"): + compare_evaluations(result(),candidate) + + def test_explicit_tolerance_is_visible_and_bounded(self): + comparison=compare_evaluations(result(success=1.0),result(success=0.99),max_task_success_drop=0.02) + self.assertTrue(comparison["gate"]["task_success_preserved"]) + + +if __name__=="__main__":unittest.main() diff --git a/test/test_radeon_forge_hooks.py b/test/test_radeon_forge_hooks.py new file mode 100644 index 0000000000000..6d86c6cc3d3d2 --- /dev/null +++ b/test/test_radeon_forge_hooks.py @@ -0,0 +1,124 @@ +import hashlib +import tempfile +import unittest +from pathlib import Path + +from extra.radeon_forge.backends.stage_hooks import ModelStageHookRuntime +from extra.radeon_forge.synthesis import SafeHookRegistry +from extra.radeon_forge.synthesis.hooks import (ExecutionContext, ExecutionStage, HookActivationError, HookRegistry, + RuntimeFingerprint, StagePredicate) +from extra.radeon_forge.synthesis.workspace import CandidateWorkspace, KernelSpec + + +class FakeModel: + def __init__(self): + self.layers = [lambda value: value + 1, lambda value: value + 2] + self.forward_jit = None + + +class TestStageAwareHooks(unittest.TestCase): + def test_predicate_distinguishes_execution_state_and_workload(self): + predicate = StagePredicate.from_mapping({"stages": ["decode"], "batch_sizes": [1], "min_context_tokens": 1024, + "min_generated_token_index": 1, "prefix_cache": "hit", + "conditions": {"resume_after_tool": True}}) + matching = ExecutionContext(ExecutionStage.DECODE, context_tokens=2048, generated_token_index=4, + prefix_reused_tokens=512, warm=True, attributes={"resume_after_tool": True}) + self.assertTrue(predicate.matches(matching)) + self.assertFalse(predicate.matches(ExecutionContext(ExecutionStage.PREFILL, context_tokens=2048, + prefix_reused_tokens=512, attributes={"resume_after_tool": True}))) + self.assertFalse(predicate.matches(ExecutionContext(ExecutionStage.DECODE, context_tokens=512, + generated_token_index=4, prefix_reused_tokens=256, + attributes={"resume_after_tool": True}))) + self.assertFalse(predicate.matches(ExecutionContext(ExecutionStage.DECODE, context_tokens=2048, + generated_token_index=0, prefix_reused_tokens=256, + attributes={"resume_after_tool": True}))) + + def test_prefill_and_decode_implementations_can_coexist_in_generic_resolver(self): + with tempfile.TemporaryDirectory() as directory: + workspace = CandidateWorkspace(directory) + common = {"layer": "transformer_block", "target": "llama.block", "adapter": "python_transformer_block", + "exclusive_group": "llama-block"} + prefill_spec = workspace.save_spec(KernelSpec("prefill-block", "prefill block", target="gfx1100", + invariants=("match reference",), metadata={"hook": {**common, "when": {"stages": ["prefill"]}}})) + decode_spec = workspace.save_spec(KernelSpec("decode-block", "decode block", target="gfx1100", + invariants=("match reference",), metadata={"hook": {**common, "when": {"stages": ["decode"], + "min_generated_token_index": 1}}})) + prefill = workspace.create_candidate(prefill_spec.spec_id, "def build_replacement(*args): return args[0]\n", "prefill") + decode = workspace.create_candidate(decode_spec.spec_id, "def build_replacement(*args): return args[0]\n", "decode") + workspace.update(prefill.candidate_id, "hardware_passed", {}) + workspace.update(decode.candidate_id, "hardware_passed", {}) + registry = HookRegistry(workspace) + fingerprint = RuntimeFingerprint(architecture="gfx1100", runtime="tinygrad", model_family="llama") + registry.activate(prefill.candidate_id, fingerprint, "validated prefill path") + registry.activate(decode.candidate_id, fingerprint, "validated decode path") + self.assertEqual(len(registry.active()), 2) + self.assertEqual(registry.resolve(ExecutionContext(ExecutionStage.PREFILL))[0].candidate_id, prefill.candidate_id) + self.assertEqual(registry.resolve(ExecutionContext(ExecutionStage.FIRST_TOKEN, generated_token_index=0)), []) + self.assertEqual(registry.resolve(ExecutionContext(ExecutionStage.DECODE, generated_token_index=2))[0].candidate_id, + decode.candidate_id) + + def test_engine_registry_rejects_unimplemented_metadata_adapter(self): + with tempfile.TemporaryDirectory() as directory: + workspace = CandidateWorkspace(directory) + spec = workspace.save_spec(KernelSpec("advisory-kernel", "kernel idea", target="gfx1100", + invariants=("match reference",), metadata={"hook":{"layer":"kernel", "target":"gemv", + "adapter":"request_metadata", "when":{"stages":["decode"]}}})) + candidate = workspace.create_candidate(spec.spec_id, "# advisory only\n", "prior kernel schedule") + workspace.update(candidate.candidate_id, "hardware_passed", {}) + registry = SafeHookRegistry(workspace) + with self.assertRaisesRegex(HookActivationError, "not executable"): + registry.activate(candidate.candidate_id, RuntimeFingerprint(architecture="gfx1100"), "deploy") + + def test_stateful_hook_requires_phase_transition_oracle(self): + with tempfile.TemporaryDirectory() as directory: + workspace = CandidateWorkspace(directory) + hook = {"layer":"transformer_block", "target":"llama.block", "adapter":"python_transformer_block", + "when":{"stages":["decode"]}} + no_transition = workspace.save_spec(KernelSpec("no-transition", "decode block", target="gfx1100", + invariants=("preserve KV state",), metadata={"hook":hook})) + first = workspace.create_candidate(no_transition.spec_id, "def build_replacement(*args): return args[0]\n", "decode") + workspace.update(first.candidate_id, "hardware_passed", {}) + registry = SafeHookRegistry(workspace) + fingerprint = RuntimeFingerprint(architecture="gfx1100", runtime="tinygrad", model_family="llama") + with self.assertRaisesRegex(HookActivationError, "phase-transition oracle"): + registry.activate(first.candidate_id, fingerprint, "deploy") + + with_transition = workspace.save_spec(KernelSpec("with-transition", "decode block", target="gfx1100", + invariants=("preserve KV state",), metadata={"hook":hook, + "heldout_command":["python3", "validate_prefill_decode_transition.py"]})) + second = workspace.create_candidate(with_transition.spec_id, "def build_replacement(*args): return args[0]\n", "decode") + workspace.update(second.candidate_id, "hardware_passed", {}) + with self.assertRaisesRegex(HookActivationError, "has not passed"): + registry.activate(second.candidate_id, fingerprint, "deploy") + workspace.update(second.candidate_id, "heldout_passed", {"heldout":{"transition":"prefill->first_token->decode"}}) + active = registry.activate(second.candidate_id, fingerprint, "transition oracle passed") + self.assertEqual(active.candidate_id, second.candidate_id) + + def test_model_adapter_rolls_back_bad_stage_hook(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + good = root / "good.py" + good.write_text("def build_replacement(original, layer_index, model, context):\n return lambda value: original(value) * 10\n", + encoding="utf-8") + bad = root / "bad.py" + bad.write_text("VALUE = 1\n", encoding="utf-8") + descriptor = {"layer": "transformer_block", "target": "llama.block", "mode": "replace", + "adapter": "python_transformer_block", "selector": {"indices": [0]}, + "when": {"stages": ["decode"]}, "priority": 0, "exclusive_group": "block", "description": ""} + def hook(path, activation): + return {"activation_id": activation, "candidate_id": activation, "spec_id": activation, + "source_path": str(path), "source_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + "descriptor": descriptor} + + model = FakeModel() + runtime = ModelStageHookRuntime(model) + baseline = model.layers[0](2) + applied = runtime.apply([hook(good, "good")], ExecutionContext(ExecutionStage.DECODE, generated_token_index=2)) + self.assertEqual(applied["active"], ["good"]) + self.assertEqual(model.layers[0](2), baseline * 10) + rolled_back = runtime.apply([hook(bad, "bad")], ExecutionContext(ExecutionStage.DECODE, generated_token_index=2)) + self.assertTrue(rolled_back["rolled_back"]) + self.assertEqual(model.layers[0](2), baseline) + + +if __name__ == "__main__": unittest.main() diff --git a/test/test_radeon_forge_kv.py b/test/test_radeon_forge_kv.py new file mode 100644 index 0000000000000..2883eafc12bde --- /dev/null +++ b/test/test_radeon_forge_kv.py @@ -0,0 +1,35 @@ +import unittest + +from extra.radeon_forge.backends.kv_state import KVReuseLedger + + +class TestKVReuseLedger(unittest.TestCase): + def test_emitted_but_unprocessed_token_is_not_reported_as_cache_hit(self): + ledger = KVReuseLedger() + prompt = [1, 2, 3] + self.assertEqual(ledger.begin("session-a", prompt), 0) + ledger.commit_prefill(prompt[:-1]) + ledger.commit_decode_input(prompt[-1], expected_position=2) + self.assertEqual(ledger.cached_tokens, [1, 2, 3]) + + # Token 4 was emitted to the user but generation stopped before it became + # the input to another decode step. It is intentionally absent from KV. + next_prompt = [1, 2, 3, 4, 5] + self.assertEqual(ledger.begin("session-a", next_prompt), 3) + ledger.commit_prefill(next_prompt[:-1]) + self.assertEqual(ledger.cached_tokens, [1, 2, 3, 4]) + + def test_session_switch_invalidates_single_resident_cache(self): + ledger = KVReuseLedger("session-a", [1, 2, 3]) + self.assertEqual(ledger.begin("session-b", [1, 2, 3]), 0) + self.assertEqual(ledger.active_session, "session-b") + self.assertEqual(ledger.cached_tokens, []) + + def test_decode_position_mismatch_fails_closed(self): + ledger = KVReuseLedger("session-a", [1, 2]) + with self.assertRaisesRegex(RuntimeError, "position mismatch"): + ledger.commit_decode_input(3, expected_position=4) + self.assertEqual(ledger.cached_tokens, [1, 2]) + + +if __name__ == "__main__": unittest.main() diff --git a/test/test_radeon_forge_openai.py b/test/test_radeon_forge_openai.py new file mode 100644 index 0000000000000..f91ed09315afd --- /dev/null +++ b/test/test_radeon_forge_openai.py @@ -0,0 +1,74 @@ +import json +import tempfile +import unittest + +from extra.radeon_forge.backends.fake import ScriptedBackend +from extra.radeon_forge.runtime import BackendCapabilities, ForgeEngine, GenerationEvent +from extra.radeon_forge.ui.openai_api import OpenAIRequestError, collect_chat_completion, stream_chat_completion + + +class ToolBackend: + RAW = '{"id":"call-1","name":"read_file","arguments":{"path":"README.md"}}' + @property + def name(self): return "tool-model" + @property + def capabilities(self): return BackendCapabilities(structured_tools=True, local_only=True) + @property + def runtime_metadata(self): return {"architecture":"gfx1100", "runtime":"tinygrad", "model_family":"llama"} + def stream(self, request): + yield GenerationEvent("prefill", metrics={"prompt_tokens":12, "wall_ms":2.0}) + yield GenerationEvent("token", self.RAW[:20], metrics={"stage":"first_token", "wall_ms":1.0}) + yield GenerationEvent("token", self.RAW[20:], metrics={"stage":"decode", "wall_ms":1.0}) + yield GenerationEvent("tool_call", tool_call={"id":"call-1", "name":"read_file", + "arguments":{"path":"README.md"}, "raw":self.RAW}) + yield GenerationEvent("done", finish_reason="tool_call", metrics={"generated_tokens":2}) + def close(self): pass + + +class TestOpenAIAdapter(unittest.TestCase): + def test_regular_completion_has_standard_shape_and_forge_metrics(self): + with tempfile.TemporaryDirectory() as directory: + engine = ForgeEngine(ScriptedBackend(["hello local world"]), directory) + result = collect_chat_completion(engine, {"model":"local", "messages":[{"role":"user", "content":"hello"}], + "user":"stable-session"}) + self.assertEqual(result["object"], "chat.completion") + self.assertEqual(result["choices"][0]["message"]["content"], "hello local world") + self.assertEqual(result["choices"][0]["finish_reason"], "stop") + self.assertEqual(result["forge"]["session_id"], "stable-session") + self.assertGreater(result["usage"]["completion_tokens"], 0) + engine.close() + + def test_tool_completion_is_structured_and_raw_protocol_is_not_content(self): + with tempfile.TemporaryDirectory() as directory: + engine = ForgeEngine(ToolBackend(), directory) + payload = {"messages":[{"role":"user", "content":"read README"}], "tools":[{"type":"function", + "function":{"name":"read_file", "description":"read", "parameters":{"type":"object"}}}]} + result = collect_chat_completion(engine, payload) + message = result["choices"][0]["message"] + self.assertIsNone(message["content"]) + self.assertEqual(message["tool_calls"][0]["function"]["name"], "read_file") + self.assertEqual(json.loads(message["tool_calls"][0]["function"]["arguments"]), {"path":"README.md"}) + self.assertEqual(result["choices"][0]["finish_reason"], "tool_calls") + engine.close() + + def test_stream_suppresses_internal_tool_markup(self): + with tempfile.TemporaryDirectory() as directory: + engine = ForgeEngine(ToolBackend(), directory) + payload = {"stream":True, "messages":[{"role":"user", "content":"read README"}], "tools":[{"type":"function", + "function":{"name":"read_file", "description":"read", "parameters":{"type":"object"}}}]} + stream = list(stream_chat_completion(engine, payload)) + joined = "".join(stream) + self.assertNotIn("", joined) + self.assertIn('"tool_calls"', joined) + self.assertIn('"finish_reason":"tool_calls"', joined) + self.assertEqual(stream[-1], "data: [DONE]\n\n") + engine.close() + + def test_invalid_messages_fail_before_model_execution(self): + with tempfile.TemporaryDirectory() as directory: + engine = ForgeEngine(ScriptedBackend(), directory) + with self.assertRaises(OpenAIRequestError): collect_chat_completion(engine, {"messages":[]}) + engine.close() + + +if __name__ == "__main__": unittest.main() diff --git a/test/test_radeon_forge_persistence.py b/test/test_radeon_forge_persistence.py new file mode 100644 index 0000000000000..99020197d894b --- /dev/null +++ b/test/test_radeon_forge_persistence.py @@ -0,0 +1,79 @@ +import tempfile +import unittest +from pathlib import Path + +from extra.radeon_forge.backends.fake import ScriptedBackend +from extra.radeon_forge.runtime import ForgeEngine +from extra.radeon_forge.runtime.session import AgentSession, SessionState +from extra.radeon_forge.runtime.store import SessionStore +from extra.radeon_forge.runtime.tools import ToolRegistry +from extra.radeon_forge.permissions import PermissionController + + +class TestSessionPersistence(unittest.TestCase): + def test_completed_session_restores_and_continues(self): + with tempfile.TemporaryDirectory() as directory: + first=ForgeEngine(ScriptedBackend(["first local answer"]),directory) + session=first.create_session(); session_id=session.session_id + first.run_message(session_id,"first question") + trace_id=session.trace.trace_id + first.close() + + second=ForgeEngine(ScriptedBackend(["continued after restart"]),directory) + restored=second.session(session_id) + self.assertEqual(restored.state,SessionState.COMPLETED) + self.assertEqual(restored.trace.trace_id,trace_id) + self.assertEqual(restored.messages[-1]["content"],"first local answer") + second.run_message(session_id,"continue") + self.assertEqual(restored.messages[-1]["content"],"continued after restart") + self.assertGreater(len(restored.trace.events()),0) + second.close() + + def test_pending_tool_approval_restores_without_execution(self): + with tempfile.TemporaryDirectory() as directory: + Path(directory,"hello.txt").write_text("private",encoding="utf-8") + first=ForgeEngine(ScriptedBackend(['{"name":"read_file","arguments":{"path":"hello.txt"}}']),directory) + session=first.create_session(); session_id=session.session_id + first.run_message(session_id,"read the file") + self.assertEqual(session.state,SessionState.AWAITING_TOOL_APPROVAL) + first.close() + + second=ForgeEngine(ScriptedBackend(["read completed after approval"]),directory) + restored=second.session(session_id) + self.assertEqual(restored.state,SessionState.AWAITING_TOOL_APPROVAL) + self.assertEqual(restored.pending_tool_call.name,"read_file") + self.assertFalse(any(event.kind=="tool_started" for event in restored.events)) + second.run_tool_approval(session_id,"approve restored local read") + self.assertEqual(restored.state,SessionState.COMPLETED) + self.assertEqual(restored.messages[-1]["content"],"read completed after approval") + second.close() + + def test_inflight_checkpoint_recovers_to_idle_without_replay(self): + with tempfile.TemporaryDirectory() as directory: + root=Path(directory)/"sessions" + store=SessionStore(root) + backend=ScriptedBackend(["unused"]) + session=AgentSession(backend,ToolRegistry(PermissionController()),"system") + session.messages.append({"role":"user","content":"incomplete turn"}) + session.state=SessionState.GENERATING + session._emit("generation_started",backend="fake",step=1,capabilities={}) + store.save(session) + restored=store.restore(session.session_id,backend,ToolRegistry(PermissionController()),"system") + self.assertEqual(restored.state,SessionState.IDLE) + recovery=next(event for event in restored.events if event.kind=="session_recovered") + self.assertEqual(recovery.data["prior_state"],"generating") + self.assertFalse(any(message.get("role")=="assistant" for message in restored.messages)) + + def test_checkpoint_path_and_format_are_local_and_atomic(self): + with tempfile.TemporaryDirectory() as directory: + store=SessionStore(directory) + session=AgentSession(ScriptedBackend(),ToolRegistry(PermissionController()),"system") + path=store.save(session) + self.assertEqual(path.parent.resolve(),Path(directory).resolve()) + self.assertTrue(path.name.endswith(".json")) + self.assertFalse(list(Path(directory).glob("*.tmp"))) + payload=store.load_payload(session.session_id) + self.assertEqual(payload["format_version"],SessionStore.FORMAT_VERSION) + + +if __name__=="__main__": unittest.main() diff --git a/test/test_radeon_forge_plugins.py b/test/test_radeon_forge_plugins.py new file mode 100644 index 0000000000000..3dadf34a0fe45 --- /dev/null +++ b/test/test_radeon_forge_plugins.py @@ -0,0 +1,41 @@ +import sys +import unittest +from pathlib import Path + +from extra.radeon_forge.backends.plugins import WorkerPlugin, WorkerPluginRegistry, default_worker_plugins + + +class TestWorkerPlugins(unittest.TestCase): + def test_legacy_llama_command_is_local_and_explicit(self): + plugin=default_worker_plugins().get("legacy-llama") + command=plugin.command({"model":Path("/models/local.gguf"),"tokenizer":Path("/models/tokenizer.model"), + "size":"1B","max_context":4096,"seed":7}) + self.assertEqual(command[:3],[sys.executable,"-m","extra.radeon_forge.backends.tinygrad_llama_worker"]) + self.assertIn("/models/local.gguf",command) + self.assertIn("--max-context",command) + self.assertIn("4096",command) + self.assertNotIn("http", " ".join(command)) + + def test_required_and_unknown_fields_fail_closed(self): + plugin=default_worker_plugins().get("legacy-llama") + with self.assertRaisesRegex(ValueError,"missing required"): + plugin.command({}) + with self.assertRaisesRegex(ValueError,"unsupported fields"): + plugin.command({"model":"x","remote_api":"https://example.com"}) + + def test_registry_rejects_external_modules_and_duplicates(self): + registry=WorkerPluginRegistry() + with self.assertRaisesRegex(ValueError,"in-tree local"): + registry.register(WorkerPlugin("remote","third_party.worker","bad",("model",))) + local=WorkerPlugin("local","extra.radeon_forge.backends.worker","ok",("model",)) + registry.register(local) + with self.assertRaisesRegex(ValueError,"duplicate"): + registry.register(local) + + def test_discovery_is_stable(self): + rows=default_worker_plugins().list() + self.assertEqual([row["name"] for row in rows],["legacy-llama"]) + self.assertIn("stage hooks",rows[0]["description"]) + + +if __name__=="__main__": unittest.main() diff --git a/test/test_radeon_forge_profile_compare.py b/test/test_radeon_forge_profile_compare.py new file mode 100644 index 0000000000000..ea936b374ec1c --- /dev/null +++ b/test/test_radeon_forge_profile_compare.py @@ -0,0 +1,81 @@ +import csv +import hashlib +import json +import tempfile +import unittest +from pathlib import Path + +from extra.radeon_forge.profiling.compare import IncomparableCaptures, compare_captures + + +class TestProfileCaptureComparison(unittest.TestCase): + def _capture(self, root: Path, capture_id: str, stage: str, rows: list[dict[str, str]], *, command=None, + counters=("SQ_WAVES", "DurationNs"), passed=True, workload_hash="suite-1") -> Path: + directory = root / capture_id + directory.mkdir() + csv_path = directory / "counters.csv" + with csv_path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=["Kernel_Name", "SQ_WAVES", "DurationNs"]) + writer.writeheader(); writer.writerows(rows) + data = csv_path.read_bytes() + numeric = {} + for key in ("SQ_WAVES", "DurationNs"): + values = [float(row[key]) for row in rows] + numeric[key] = {"count":len(values), "min":min(values), "max":max(values), "mean":sum(values)/len(values)} + payload = { + "capture_id":capture_id, "kind":"counters", "passed":passed, + "target_command":command or ["python3", "workload.py", "--case", "decode"], + "artifacts":[{"path":str(csv_path), "size_bytes":len(data), "sha256":hashlib.sha256(data).hexdigest(), "suffix":".csv"}], + "summary":{"stage":stage, "counters":list(counters), "capture_metadata":{"workload_hash":workload_hash}, + "parsed_csv":{"csv_files":1, "csv_rows":len(rows), "numeric_columns":numeric}} + } + manifest = directory / "forge_capture_manifest.json" + manifest.write_text(json.dumps(payload), encoding="utf-8") + return manifest + + def test_compatible_captures_produce_global_and_per_kernel_deltas(self): + with tempfile.TemporaryDirectory() as directory: + root=Path(directory) + baseline=self._capture(root,"base","decode",[ + {"Kernel_Name":"decode_gemv","SQ_WAVES":"64","DurationNs":"1000"}, + {"Kernel_Name":"decode_gemv","SQ_WAVES":"64","DurationNs":"1200"}, + ]) + candidate=self._capture(root,"candidate","decode",[ + {"Kernel_Name":"decode_gemv","SQ_WAVES":"72","DurationNs":"800"}, + {"Kernel_Name":"decode_gemv","SQ_WAVES":"72","DurationNs":"900"}, + ]) + result=compare_captures(baseline,candidate) + self.assertTrue(result["compatibility"]["comparable"]) + self.assertEqual(result["numeric_column_deltas"]["DurationNs"]["baseline"],1100.0) + self.assertEqual(result["numeric_column_deltas"]["DurationNs"]["candidate"],850.0) + row=next(item for item in result["per_kernel_deltas"] if item["kernel"]=="decode_gemv" and item["metric"]=="DurationNs") + self.assertEqual(row["absolute"],-250.0) + self.assertAlmostEqual(row["relative"],-250/1100) + self.assertEqual(result["interpretation_contract"]["lower_is_better_metrics"],[]) + + def test_stage_or_workload_mismatch_is_rejected(self): + with tempfile.TemporaryDirectory() as directory: + root=Path(directory) + rows=[{"Kernel_Name":"k","SQ_WAVES":"1","DurationNs":"1"}] + baseline=self._capture(root,"base","prefill",rows) + candidate=self._capture(root,"candidate","decode",rows) + with self.assertRaisesRegex(IncomparableCaptures,"stage differs"): + compare_captures(baseline,candidate) + other=self._capture(root,"other","prefill",rows,workload_hash="suite-2") + with self.assertRaisesRegex(IncomparableCaptures,"workload_identity differs"): + compare_captures(baseline,other) + + def test_counter_set_or_failed_capture_is_rejected(self): + with tempfile.TemporaryDirectory() as directory: + root=Path(directory) + rows=[{"Kernel_Name":"k","SQ_WAVES":"1","DurationNs":"1"}] + baseline=self._capture(root,"base","decode",rows) + different=self._capture(root,"different","decode",rows,counters=("SQ_WAVES",)) + with self.assertRaisesRegex(IncomparableCaptures,"counters differs"): + compare_captures(baseline,different) + failed=self._capture(root,"failed","decode",rows,passed=False) + with self.assertRaisesRegex(IncomparableCaptures,"candidate capture did not pass"): + compare_captures(baseline,failed) + + +if __name__=="__main__": unittest.main() diff --git a/test/test_radeon_forge_profile_evidence.py b/test/test_radeon_forge_profile_evidence.py new file mode 100644 index 0000000000000..85711716e9dcb --- /dev/null +++ b/test/test_radeon_forge_profile_evidence.py @@ -0,0 +1,39 @@ +import unittest + +from extra.radeon_forge.profiling.report import build_profile_report +from extra.radeon_forge.runtime.events import TraceRecorder + + +class TestHardwareProfileEvidence(unittest.TestCase): + def test_att_and_counter_captures_are_observed_evidence(self): + trace=TraceRecorder("trace-1") + trace.point("profile","hardware_capture",capture_id="att-1",capture_kind="att",passed=True,stage="decode", + artifact_count=3,output_directory="/tmp/att",parsed_csv={},error="") + trace.point("profile","hardware_capture",capture_id="counter-1",capture_kind="counters",passed=True,stage="first_token", + artifact_count=2,output_directory="/tmp/counters",parsed_csv={"csv_rows":8},error="") + report=build_profile_report(trace.events()) + self.assertEqual(report["summary"]["hardware_capture_count"],2) + self.assertEqual(report["summary"]["att_capture_count"],1) + self.assertEqual(report["summary"]["counter_capture_count"],1) + titles={finding["title"] for finding in report["findings"]} + self.assertIn("AMD ATT/SQTT evidence is attached",titles) + self.assertIn("ROCm hardware counter evidence is attached",titles) + self.assertNotIn("No ROCm counter or ATT capture is attached to this trace",titles) + + def test_failed_capture_is_visible(self): + trace=TraceRecorder("trace-2") + trace.point("profile","hardware_capture",capture_id="att-bad",capture_kind="att",passed=False,stage="decode", + artifact_count=0,output_directory="/tmp/att-bad",parsed_csv={},error="rocprofv3 missing ATT support") + report=build_profile_report(trace.events()) + self.assertEqual(report["summary"]["failed_capture_count"],1) + finding=next(item for item in report["findings"] if item["title"]=="One or more hardware profiling captures failed") + self.assertEqual(finding["status"],"observed") + self.assertIn("missing ATT support",finding["evidence"][0]) + + def test_absence_remains_unknown(self): + report=build_profile_report(()) + finding=next(item for item in report["findings"] if item["title"]=="No ROCm counter or ATT capture is attached to this trace") + self.assertEqual(finding["status"],"unknown") + + +if __name__=="__main__": unittest.main() diff --git a/test/test_radeon_forge_recipe.py b/test/test_radeon_forge_recipe.py new file mode 100644 index 0000000000000..cadebef21c5fd --- /dev/null +++ b/test/test_radeon_forge_recipe.py @@ -0,0 +1,155 @@ +import base64 +import tempfile +import unittest +from pathlib import Path + +from extra.radeon_forge.synthesis import (CandidateWorkspace, ForgeRecipe, HookDescriptor, KernelSpec, RecipeLibrary, + export_recipe, export_recipe_with_hook) + + +class TestForgeRecipe(unittest.TestCase): + def _recipe_text(self, seed: str = "print('seed')\n") -> str: + encoded = base64.b64encode(seed.encode()).decode() + return f'''format_version = 1 +name = "portable-vopd" +description = "One-file optimization knowledge" +operation = "Tune an RDNA3 instruction schedule" +target = "gfx1100" +objective = "minimize P95 latency" +agent_brief = "Inspect evidence, freely regenerate the implementation, and trust only the oracle." +invariants = ["match reference", "never use MockGPU timing"] +unknowns = ["best schedule is hardware dependent"] +seed_artifact = "seed.py" + +[compatibility] +architecture = "gfx1100" + +[acceptance] +maximum_error = 0.000001 + +[metadata.hook] +layer = "kernel" +target = "projection_gemm" +mode = "replace" +adapter = "request_metadata" +exclusive_group = "projection" + +[metadata.hook.when] +stages = ["decode"] +batch_sizes = [1] +min_context_tokens = 1024 +prefix_cache = "hit" + +[metadata.search] +budgets = [3, 10] +reduction = 2 +max_candidates = 16 + +[metadata.search.axes] +WAVES = [2, 4] +STAGES = [1, 2] + +[oracle] +mockgpu_command = ["python3", "{{candidate}}", "--bundle", "{{bundle}}"] +hardware_command = ["python3", "{{candidate}}", "--benchmark"] + +[[artifact]] +path = "seed.py" +role = "implementation_cache" +sha256 = "{__import__('hashlib').sha256(seed.encode()).hexdigest()}" +content_base64 = "{encoded}" + +[[artifact]] +path = "knowledge/notes.md" +role = "knowledge" +content = "The implementation is disposable." +''' + + def test_install_is_content_addressed_and_seed_is_unverified(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source_a, source_b = root / "a.forge.toml", root / "b.forge.toml" + source_a.write_text(self._recipe_text(), encoding="utf-8") + source_b.write_text(self._recipe_text(), encoding="utf-8") + recipe_a, recipe_b = ForgeRecipe.load(source_a), ForgeRecipe.load(source_b) + self.assertEqual(recipe_a.recipe_id, recipe_b.recipe_id) + + workspace = CandidateWorkspace(root / "library") + installed = RecipeLibrary(workspace).install(recipe_a) + self.assertTrue((Path(installed.bundle_root) / "knowledge/notes.md").is_file()) + self.assertIsNotNone(installed.seed_candidate_id) + candidate = workspace.load_candidate(installed.seed_candidate_id) + self.assertEqual(candidate.status, "imported_unverified") + self.assertEqual(Path(candidate.source_path).read_text(encoding="utf-8"), "print('seed')\n") + + spec = workspace.load_spec(installed.spec_id) + descriptor = HookDescriptor.from_spec(spec) + self.assertEqual(descriptor.layer.value, "kernel") + self.assertEqual([x.value for x in descriptor.when.stages], ["decode"]) + self.assertEqual(descriptor.when.min_context_tokens, 1024) + self.assertEqual(descriptor.when.prefix_cache, "hit") + self.assertEqual(spec.metadata["search"]["axes"]["WAVES"], [2, 4]) + rendered = workspace.render_command(spec.mockgpu_command, spec, candidate, root) + self.assertEqual(rendered[1], candidate.source_path) + self.assertEqual(rendered[3], installed.bundle_root) + + def test_rejects_path_traversal(self): + text = self._recipe_text().replace('path = "seed.py"', 'path = "../seed.py"', 1).replace('seed_artifact = "seed.py"', 'seed_artifact = "../seed.py"') + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "bad.forge.toml" + path.write_text(text, encoding="utf-8") + with self.assertRaises(ValueError): ForgeRecipe.load(path) + + def test_export_import_round_trip_preserves_contract_cache_stage_and_search(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source_workspace = CandidateWorkspace(root / "source") + spec = KernelSpec( + name="roundtrip", operation="specialize batch-one decode", target="gfx1100", language="python", extension=".py", + invariants=("match reference", "no remote API"), objective="minimize p95 token latency", + mockgpu_command=("python3", "{candidate}"), hardware_command=("python3", "{candidate}", "--benchmark"), + metadata={"recipe_agent_brief":"Use the oracle as the contract and freely rewrite the implementation.", + "recipe_compatibility":{"architecture":"gfx1100"}, "recipe_acceptance":{"max_error":1e-6}, + "search":{"axes":{"WAVES":[2,4], "UNROLL":[1,2]}, "budgets":[3,10,30], "reduction":2, + "max_candidates":16}, + "hook":{"layer":"transformer_block", "target":"llama.decode.block", "mode":"replace", + "adapter":"python_transformer_block", "selector":{"indices":[0,1]}, + "when":{"stages":["decode"], "batch_sizes":[1], "min_generated_token_index":1, + "conditions":{"resume_after_tool":False}}, + "exclusive_group":"decode-block"}}, + ) + source_workspace.save_spec(spec) + candidate = source_workspace.create_candidate(spec.spec_id, "print('candidate')\n", "measured decode implementation") + source_workspace.update(candidate.candidate_id, "hardware_passed", {"selected_parameters":{"WAVES":4, "UNROLL":2}}) + exported = export_recipe_with_hook(source_workspace, spec.spec_id, root / "shared.forge.toml", candidate.candidate_id) + + loaded = ForgeRecipe.load(exported) + self.assertEqual(loaded.target, "gfx1100") + self.assertEqual(loaded.invariants, spec.invariants) + self.assertEqual(loaded.metadata["hook"]["when"]["stages"], ["decode"]) + self.assertEqual(loaded.metadata["search"]["axes"]["WAVES"], [2, 4]) + self.assertEqual(loaded.metadata["exported_winner"], {"WAVES":4, "UNROLL":2}) + destination = CandidateWorkspace(root / "destination") + installed = RecipeLibrary(destination).install(loaded) + imported = destination.load_candidate(installed.seed_candidate_id) + self.assertEqual(Path(imported.source_path).read_text(encoding="utf-8"), "print('candidate')\n") + self.assertEqual(imported.status, "imported_unverified") + installed_spec = destination.load_spec(installed.spec_id) + descriptor = HookDescriptor.from_spec(installed_spec) + self.assertEqual(descriptor.target, "llama.decode.block") + self.assertEqual([x.value for x in descriptor.when.stages], ["decode"]) + self.assertEqual(descriptor.when.min_generated_token_index, 1) + self.assertEqual(descriptor.when.conditions["resume_after_tool"], False) + self.assertEqual(installed_spec.metadata["search"]["axes"]["UNROLL"], [1, 2]) + self.assertEqual(installed_spec.metadata["exported_winner"]["WAVES"], 4) + + def test_base_export_remains_valid_without_hook_extension(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + workspace = CandidateWorkspace(root / "source") + spec = workspace.save_spec(KernelSpec("plain", "plain operation", invariants=("correct",))) + path = export_recipe(workspace, spec.spec_id, root / "plain.forge.toml") + self.assertEqual(ForgeRecipe.load(path).name, "plain") + + +if __name__ == "__main__": unittest.main() diff --git a/test/test_radeon_forge_runtime.py b/test/test_radeon_forge_runtime.py new file mode 100644 index 0000000000000..ce1ed413ef422 --- /dev/null +++ b/test/test_radeon_forge_runtime.py @@ -0,0 +1,176 @@ +import tempfile +import threading +import time +import unittest +from pathlib import Path + +from extra.radeon_forge.backends.fake import ScriptedBackend +from extra.radeon_forge.permissions import PermissionDenied +from extra.radeon_forge.profiling.report import build_profile_report +from extra.radeon_forge.runtime import ForgeEngine +from extra.radeon_forge.runtime.backend import BackendCapabilities, GenerationEvent +from extra.radeon_forge.runtime.session import SessionState +from extra.radeon_forge.runtime.tools import WorkspaceTools +from extra.radeon_forge.synthesis.workspace import CandidateWorkspace, KernelSpec + + +class KernelEvidenceBackend: + @property + def name(self): return "kernel-evidence-local" + @property + def capabilities(self): return BackendCapabilities(prefix_cache=True, persistent_kv=True, kernel_metrics=True) + @property + def runtime_metadata(self): return {"runtime":"test", "architecture":"gfx1100", "model_family":"llama"} + def stream(self, request): + yield GenerationEvent("prefill", metrics={"wall_ms": 8.0, "gpu_ms": 5.0, "prompt_tokens": 24, + "prefix_reused_tokens": 12, "profile_kernel_events": 1}) + yield GenerationEvent("kernel", metrics={"name": "rmsnorm_fused", "duration_ms": 0.40, "stage": "prefill", "device": "AMD"}) + yield GenerationEvent("token", "first", metrics={"index": 0, "stage": "first_token", "wall_ms": 4.0, "gpu_ms": 2.5, + "kernel_count": 24, "profile_kernel_events": 1}) + yield GenerationEvent("kernel", metrics={"name": "first_token_projection", "duration_ms": 0.80, + "stage": "first_token", "device": "AMD"}) + yield GenerationEvent("token", "ok", metrics={"index": 1, "stage": "decode", "wall_ms": 2.0, "gpu_ms": 1.4, + "kernel_count": 24, "profile_kernel_events": 2}) + yield GenerationEvent("kernel", metrics={"name": "decode_gemv", "duration_ms": 0.90, "stage": "decode", "device": "AMD"}) + yield GenerationEvent("kernel", metrics={"name": "decode_gemv", "duration_ms": 0.70, "stage": "decode", "device": "AMD"}) + yield GenerationEvent("done", finish_reason="stop", metrics={"generated_tokens": 2}) + def close(self): pass + + +class BlockingBackend: + def __init__(self): self.first_token = threading.Event(); self.release = threading.Event() + @property + def name(self): return "blocking-local" + @property + def capabilities(self): return BackendCapabilities(streaming=True) + @property + def runtime_metadata(self): return {"runtime":"test", "architecture":"gfx1100", "model_family":"llama"} + def stream(self, request): + yield GenerationEvent("prefill", metrics={"wall_ms": 1.0, "prompt_tokens": 8, "prefix_reused_tokens": 0}) + yield GenerationEvent("token", "partial ", metrics={"index":0, "stage":"first_token", "wall_ms":1.0}) + self.first_token.set() + if not self.release.wait(2): raise TimeoutError("test did not release backend") + yield GenerationEvent("token", "complete", metrics={"index":1, "stage":"decode", "wall_ms":1.0}) + yield GenerationEvent("done", finish_reason="stop", metrics={"generated_tokens":2}) + def close(self): self.release.set() + + +class NativeToolBackend: + RAW = '{"id":"call-7","name":"read_file","arguments":{"path":"hello.txt","start_line":2}}' + @property + def name(self): return "native-tool-local" + @property + def capabilities(self): return BackendCapabilities(structured_tools=True) + @property + def runtime_metadata(self): return {"runtime":"test", "architecture":"gfx1100", "model_family":"llama"} + def stream(self, request): + yield GenerationEvent("token", self.RAW, metrics={"index":0, "stage":"first_token", "wall_ms":1.0}) + yield GenerationEvent("tool_call", tool_call={"id":"call-7", "name":"read_file", + "arguments":{"path":"hello.txt", "start_line":2}, "raw":self.RAW}) + yield GenerationEvent("done", finish_reason="tool_call", metrics={"generated_tokens":1}) + def close(self): pass + + +class TestRadeonForgeRuntime(unittest.TestCase): + def test_kernel_evidence_survives_unified_trace_and_is_ranked_by_stage(self): + with tempfile.TemporaryDirectory() as directory: + engine = ForgeEngine(KernelEvidenceBackend(), directory) + session = engine.create_session() + session.send("profile two execution states") + report = build_profile_report(session.trace.events()) + top = report["summary"]["top_kernels"] + self.assertEqual(top[0]["name"], "decode_gemv") + self.assertEqual(top[0]["calls"], 2) + self.assertAlmostEqual(top[0]["total_ms"], 1.6) + self.assertEqual(report["summary"]["top_kernels_by_stage"]["prefill"][0]["name"], "rmsnorm_fused") + self.assertEqual(report["summary"]["top_kernels_by_stage"]["first_token"][0]["name"], "first_token_projection") + self.assertEqual(report["summary"]["top_kernels_by_stage"]["decode"][0]["name"], "decode_gemv") + self.assertTrue(report["summary"]["kernel_profile_complete"]) + titles = {x["title"] for x in report["findings"]} + self.assertIn("One kernel family dominates captured GPU time overall", titles) + self.assertIn("Different execution stages have different dominant kernels", titles) + self.assertIn("First-token and steady-decode latency are materially different", titles) + kernel_events = [event for event in session.trace.events() if event.kind == "kernel"] + self.assertEqual([event.name for event in kernel_events], ["rmsnorm_fused", "first_token_projection", "decode_gemv", "decode_gemv"]) + engine.close() + + def test_async_job_exposes_partial_generation(self): + with tempfile.TemporaryDirectory() as directory: + backend = BlockingBackend() + engine = ForgeEngine(backend, directory) + session = engine.create_session() + job = engine.submit_message(session.session_id, "stream locally") + self.assertTrue(backend.first_token.wait(1)) + self.assertEqual(engine.jobs.snapshot(job["job_id"]).state, "running") + self.assertEqual(session.state, SessionState.GENERATING) + self.assertEqual(session.partial_output, "partial ") + self.assertTrue(any(event["kind"] == "token" for event in session.events_after(0))) + backend.release.set() + deadline = time.time() + 2 + while engine.jobs.snapshot(job["job_id"]).state not in {"completed", "failed"} and time.time() < deadline: time.sleep(0.01) + self.assertEqual(engine.jobs.snapshot(job["job_id"]).state, "completed") + self.assertEqual(session.state, SessionState.COMPLETED) + self.assertEqual(session.partial_output, "") + self.assertEqual(session.messages[-1]["content"], "partial complete") + engine.close() + + def test_permissioned_tool_round_trip_preserves_arguments(self): + responses = [ + '{"name":"read_file","arguments":{"path":"hello.txt"}}', + "The private file was read successfully.", + ] + with tempfile.TemporaryDirectory() as directory: + Path(directory, "hello.txt").write_text("local only", encoding="utf-8") + engine = ForgeEngine(ScriptedBackend(responses), directory) + session = engine.create_session() + session.send("read hello.txt") + self.assertEqual(session.state, SessionState.AWAITING_TOOL_APPROVAL) + with self.assertRaises(PermissionDenied): session.approve_tool("not-a-real-grant") + self.assertEqual(session.state, SessionState.AWAITING_TOOL_APPROVAL) + token = engine.grant_for_pending_tool(session.session_id, "approve one local read") + session.approve_tool(token) + self.assertEqual(session.state, SessionState.COMPLETED) + self.assertEqual(session.messages[-1]["content"], "The private file was read successfully.") + assistant_tool = next(message for message in session.messages if message.get("role") == "assistant" and "" in message.get("content", "")) + self.assertIn('"path":"hello.txt"', assistant_tool["content"]) + self.assertTrue(any(event.kind == "tool_result" for event in session.events)) + engine.close() + + def test_native_structured_tool_call_and_rejection_are_lossless(self): + with tempfile.TemporaryDirectory() as directory: + Path(directory, "hello.txt").write_text("line1\nline2\n", encoding="utf-8") + engine = ForgeEngine(NativeToolBackend(), directory) + session = engine.create_session() + session.send("read line two") + self.assertEqual(session.state, SessionState.AWAITING_TOOL_APPROVAL) + self.assertEqual(session.pending_tool_call.call_id, "call-7") + self.assertEqual(session.pending_tool_call.arguments["start_line"], 2) + self.assertEqual(session.pending_assistant_content, NativeToolBackend.RAW) + session.reject_tool("not now") + self.assertEqual(session.state, SessionState.IDLE) + self.assertEqual(session.messages[-2]["content"], NativeToolBackend.RAW) + self.assertEqual(session.messages[-1]["tool_call_id"], "call-7") + self.assertIn("not now", session.messages[-1]["content"]) + engine.close() + + def test_workspace_paths_cannot_escape(self): + with tempfile.TemporaryDirectory() as directory: + tools = WorkspaceTools(directory) + with self.assertRaisesRegex(ValueError, "escapes workspace"): + tools.read_file({"path": "../outside.txt"}) + with self.assertRaisesRegex(ValueError, "not allowlisted"): + tools.run_command({"command": "curl https://example.com"}) + + def test_generated_candidates_are_content_addressed(self): + with tempfile.TemporaryDirectory() as directory: + workspace = CandidateWorkspace(directory) + spec = workspace.save_spec(KernelSpec(name="decode-gemv", operation="gemv", shapes={"hidden": 4096}, + invariants=("numerically_matches_reference",))) + first = workspace.create_candidate(spec.spec_id, "def build_kernel():\n return 1\n", "baseline schedule") + second = workspace.create_candidate(spec.spec_id, "def build_kernel():\n return 1\n", "same implementation, new explanation") + self.assertEqual(first.candidate_id, second.candidate_id) + self.assertEqual(first.source_sha256, second.source_sha256) + self.assertTrue(Path(second.source_path).is_file()) + + +if __name__ == "__main__": unittest.main() diff --git a/test/test_radeon_forge_scaffold.py b/test/test_radeon_forge_scaffold.py new file mode 100644 index 0000000000000..f5bcff2dfdc79 --- /dev/null +++ b/test/test_radeon_forge_scaffold.py @@ -0,0 +1,38 @@ +import ast +import tempfile +import unittest +from pathlib import Path + +from extra.radeon_forge.synthesis.scaffold import render_candidate_scaffold +from extra.radeon_forge.synthesis.workspace import KernelSpec + + +class TestCandidateScaffold(unittest.TestCase): + def test_transformer_block_scaffold_is_executable_and_preserves_contract(self): + spec=KernelSpec("decode-block","specialize decode",target="gfx1100",invariants=("match reference","preserve KV"), + metadata={"search":{"axes":{"WAVES":[1,2,4]}},"hook":{"layer":"transformer_block","target":"llama.decode.block", + "mode":"replace","adapter":"python_transformer_block","selector":{"indices":"all"}, + "when":{"stages":["first_token","decode"],"batch_sizes":[1]}}}) + source=render_candidate_scaffold(spec) + ast.parse(source) + namespace={} + exec(compile(source,"candidate.py","exec"),namespace) + namespace["configure"]({"WAVES":4}) + self.assertEqual(namespace["RADEON_FORGE_PARAMETERS"],{"WAVES":4}) + original=lambda x,start_pos,freqs,mask:x+1 + replacement=namespace["build_replacement"](original,0,object(),{"stage":"decode"}) + self.assertEqual(replacement(2,0,None,None),3) + self.assertIn("first_token|decode",source) + self.assertIn("preserve KV",source) + self.assertIn('"WAVES"',source) + + def test_non_block_scaffold_does_not_invent_a_dsl(self): + spec=KernelSpec("gemv","direct gemv",invariants=("correct",),metadata={"hook":{"layer":"kernel","target":"gemv"}}) + source=render_candidate_scaffold(spec) + ast.parse(source) + self.assertIn("def build_kernel",source) + self.assertIn("NotImplementedError",source) + self.assertNotIn("class Tile",source) + + +if __name__=="__main__": unittest.main() diff --git a/test/test_radeon_forge_stop_sequences.py b/test/test_radeon_forge_stop_sequences.py new file mode 100644 index 0000000000000..9b27f2e70c7fc --- /dev/null +++ b/test/test_radeon_forge_stop_sequences.py @@ -0,0 +1,36 @@ +import unittest + +from extra.radeon_forge.runtime.stop_sequences import StopSequenceMatcher + + +class TestStopSequenceMatcher(unittest.TestCase): + def test_stop_split_across_tokens_is_not_emitted(self): + matcher = StopSequenceMatcher([""]) + self.assertEqual(matcher.feed("answerignored") + self.assertEqual(matched.text, "") + self.assertEqual(matched.matched, "") + self.assertEqual(matcher.finalize(), "") + + def test_ambiguous_prefix_is_released_when_it_stops_matching(self): + matcher = StopSequenceMatcher(["STOP"]) + self.assertEqual(matcher.feed("hello ST").text, "hello ") + self.assertEqual(matcher.feed("AR").text, "STAR") + self.assertEqual(matcher.finalize(), "") + + def test_earliest_stop_wins(self): + matcher = StopSequenceMatcher(["END", "STOP"]) + result = matcher.feed("one STOP two END") + self.assertEqual(result.text, "one ") + self.assertEqual(result.matched, "STOP") + + def test_finalize_releases_safe_suffix(self): + matcher = StopSequenceMatcher(["STOP"]) + self.assertEqual(matcher.feed("value ST").text, "value ") + self.assertEqual(matcher.finalize(), "ST") + + def test_empty_stop_is_rejected(self): + with self.assertRaises(ValueError): StopSequenceMatcher([""]) + + +if __name__ == "__main__": unittest.main() diff --git a/test/test_radeon_forge_text_events.py b/test/test_radeon_forge_text_events.py new file mode 100644 index 0000000000000..d74fcd71bb469 --- /dev/null +++ b/test/test_radeon_forge_text_events.py @@ -0,0 +1,57 @@ +import tempfile +import unittest + +from extra.radeon_forge.profiling.report import build_profile_report +from extra.radeon_forge.runtime import BackendCapabilities, ForgeEngine, GenerationEvent +from extra.radeon_forge.ui.openai_api import collect_chat_completion, stream_chat_completion + + +class BufferedTextBackend: + @property + def name(self): return "buffered-text-local" + + @property + def capabilities(self): return BackendCapabilities(streaming=True, local_only=True) + + @property + def runtime_metadata(self): return {"architecture":"gfx1100", "runtime":"tinygrad", "model_family":"llama"} + + def stream(self, request): + yield GenerationEvent("prefill", metrics={"prompt_tokens":8, "wall_ms":1.0}) + # A real model token was generated, but its characters were withheld while + # Forge determined whether they were the prefix of a cross-token stop. + yield GenerationEvent("token", "", metrics={"index":0, "stage":"first_token", "wall_ms":2.0, + "gpu_ms":1.5, "kernel_count":20}) + yield GenerationEvent("text", "visible answer", metrics={"stage":"first_token", "buffered_for_stop_sequence":True}) + yield GenerationEvent("done", finish_reason="stop", metrics={"generated_tokens":1, "stop_sequence":""}) + + def close(self): pass + + +class TestBufferedTextTransport(unittest.TestCase): + def test_agent_session_preserves_text_without_fake_token(self): + with tempfile.TemporaryDirectory() as directory: + engine = ForgeEngine(BufferedTextBackend(), directory) + session = engine.create_session() + engine.run_message(session.session_id, "answer until the stop marker") + self.assertEqual(session.messages[-1]["content"], "visible answer") + report = build_profile_report(session.trace.events()) + self.assertEqual(report["summary"]["decode_tokens"], 1) + self.assertEqual(report["summary"]["token_wall_ms_p50"], 2.0) + self.assertEqual(len([event for event in session.events if event.kind == "text"]), 1) + engine.close() + + def test_openai_completion_and_stream_include_visible_text(self): + payload = {"messages":[{"role":"user", "content":"answer"}], "stop":[""]} + with tempfile.TemporaryDirectory() as directory: + engine = ForgeEngine(BufferedTextBackend(), directory) + result = collect_chat_completion(engine, payload) + self.assertEqual(result["choices"][0]["message"]["content"], "visible answer") + self.assertEqual(result["usage"]["completion_tokens"], 1) + stream = "".join(stream_chat_completion(engine, {**payload, "session_id":"stream-session"})) + self.assertIn("visible answer", stream) + self.assertIn('"finish_reason":"stop"', stream) + engine.close() + + +if __name__ == "__main__": unittest.main() diff --git a/test/test_radeon_forge_tool_prompt.py b/test/test_radeon_forge_tool_prompt.py new file mode 100644 index 0000000000000..b581bc911dd0f --- /dev/null +++ b/test/test_radeon_forge_tool_prompt.py @@ -0,0 +1,33 @@ +import unittest + +from extra.radeon_forge.runtime.tool_prompt import TOOL_PROTOCOL_MARKER, inject_tool_instruction, render_tool_instruction + + +class TestToolPrompt(unittest.TestCase): + def test_tool_schema_is_rendered_for_local_model(self): + tools = [{"type":"function", "function":{"name":"read_file", "description":"Read a file", + "parameters":{"type":"object", "properties":{"path":{"type":"string"}}, "required":["path"]}}}] + text = render_tool_instruction(tools) + self.assertIn(TOOL_PROTOCOL_MARKER, text) + self.assertIn('"name": "read_file"', text) + self.assertIn('', text) + + def test_instruction_merges_into_existing_system_message_once(self): + tools = [{"type":"function", "function":{"name":"read_file", "parameters":{"type":"object"}}}] + original = [{"role":"system", "content":"Be precise."}, {"role":"user", "content":"Read it"}] + injected = inject_tool_instruction(original, tools) + self.assertEqual(len(injected), 2) + self.assertTrue(injected[0]["content"].startswith("Be precise.")) + self.assertEqual(injected[0]["content"].count(TOOL_PROTOCOL_MARKER), 1) + reinjected = inject_tool_instruction(injected, tools) + self.assertEqual(reinjected[0]["content"].count(TOOL_PROTOCOL_MARKER), 1) + self.assertEqual(original[0]["content"], "Be precise.") + + def test_instruction_is_prepended_when_no_system_message_exists(self): + tools = [{"type":"function", "function":{"name":"search_text", "parameters":{"type":"object"}}}] + injected = inject_tool_instruction([{"role":"user", "content":"Search"}], tools) + self.assertEqual(injected[0]["role"], "system") + self.assertIn(TOOL_PROTOCOL_MARKER, injected[0]["content"]) + + +if __name__ == "__main__": unittest.main() diff --git a/test/test_radeon_forge_tool_stream.py b/test/test_radeon_forge_tool_stream.py new file mode 100644 index 0000000000000..850154e91ee69 --- /dev/null +++ b/test/test_radeon_forge_tool_stream.py @@ -0,0 +1,35 @@ +import unittest + +from extra.radeon_forge.runtime.tool_stream import ToolProtocolError, ToolStreamParser + + +class TestToolStreamParser(unittest.TestCase): + def test_incremental_valid_tool_call(self): + parser = ToolStreamParser(["read_file", "search_text"]) + self.assertIsNone(parser.feed('{"name":"read_')) + call = parser.feed('file","arguments":{"path":"README.md"}}') + self.assertEqual(call.name, "read_file") + self.assertEqual(call.arguments, {"path":"README.md"}) + self.assertEqual(call.raw, '{"name":"read_file","arguments":{"path":"README.md"}}') + + def test_unknown_tool_fails_closed(self): + parser = ToolStreamParser(["read_file"]) + with self.assertRaisesRegex(ToolProtocolError, "unknown or unavailable"): + parser.feed('{"name":"curl","arguments":{}}') + + def test_malformed_or_incomplete_calls_fail_closed(self): + parser = ToolStreamParser(["read_file"]) + with self.assertRaisesRegex(ToolProtocolError, "exactly one"): + parser.feed('Sure. {"name":"read_file","arguments":{}}') + incomplete = ToolStreamParser(["read_file"]) + incomplete.feed('{"name":"read_file"') + with self.assertRaisesRegex(ToolProtocolError, "incomplete"): + incomplete.finalize() + + def test_stream_size_is_bounded(self): + parser = ToolStreamParser(["read_file"], max_bytes=8) + with self.assertRaisesRegex(ToolProtocolError, "size limit"): + parser.feed("123456789") + + +if __name__ == "__main__": unittest.main()