Add model-aware support for MI300X, MI325X, and MI300A

Makefile

@@ -7,7 +7,7 @@
 SHELL := /bin/bash
 
 .PHONY: help docker-shell docker-check-agents docker-smoke docker-run docker-parallel-run docker-setup-flydsl \
-        check-docker-runner check-evaluator check-held-out check-visualization \
+        check-docker-runner check-evaluator check-gpu-architecture check-held-out check-visualization \
         visualization-build visualization-serve visualization-run \
         sync-perf-helpers check-perf-helpers materialize-perf-workspace \
         materialize-perf-task cleanup-works install-cursor-agent vllm
@@ -23,12 +23,13 @@ help:
 	@echo "make docker-smoke        - Verify Docker Python, ROCm tools, imports, and GPU access"
 	@echo "make docker-run CONFIG=example_configs/quickstart_claude_mi300.yaml RUN_ARGS=\"--run-suffix test\" - Run an experiment in Docker"
 	@echo "make docker-parallel-run CONFIG=example_configs/benchmark_cursor_mi355x.yaml GPU_IDS=0,1 - Run an experiment across one worker container per GPU"
-	@echo "                         Default CONFIG is the MI300/MI300X Claude quickstart"
+	@echo "                         Default CONFIG is the generic MI300-series Claude quickstart"
 	@echo "                         On other GPUs, pass a matching CONFIG explicitly"
 	@echo "                         Images: gfx942->mi30x, gfx950->mi35x; override with AKA_DOCKER_IMAGE=..."
 	@echo "make docker-setup-flydsl - Install FlyDSL when absent (for flydsl2flydsl, torch2flydsl, and triton2flydsl)"
 	@echo "make check-docker-runner - Check Docker runner syntax and runtime-specific arguments"
 	@echo "make check-evaluator     - Run centralized evaluator unit tests"
+	@echo "make check-gpu-architecture - Check GPU model mappings and composed architecture prompts"
 	@echo "make check-held-out      - Run held-out module unit tests"
 	@echo "make visualization-run   - Build and serve the local comparison dashboard"
 	@echo "make check-visualization - Run visualization module unit tests"
@@ -80,6 +81,9 @@ check-docker-runner:
 check-evaluator:
 	@python3 -m unittest discover -s tests -p 'test_evaluator_*.py'
 
+check-gpu-architecture:
+	@python3 -m unittest discover -s tests -p 'test_gpu_architecture.py'
+
 check-held-out:
 	@python3 -m unittest discover -s tests -p 'test_held_out.py'
 

README.md

@@ -187,17 +187,22 @@ installation. The npm path requires Node.js 22+ and npm. See the
 [official Claude Code setup guide](https://code.claude.com/docs/en/installation)
 for the current alternatives.
 
-The repository provides three ready-to-use run configurations:
+The repository provides generic and model-specific quickstarts plus a longer
+MI355X benchmark configuration:
 
 | Configuration | Purpose |
 | --- | --- |
-| `example_configs/quickstart_claude_mi300.yaml` | One Claude Code GELU task on MI300/MI300X (`gfx942`); use this for a first run on MI300-series hardware. |
+| `example_configs/quickstart_claude_mi300.yaml` | Generic MI300-series (`gfx942`) Claude Code GELU task; retains backward compatibility without assuming a physical SKU. |
+| `example_configs/quickstart_claude_mi300x.yaml` | Model-specific MI300X (`gfx942`) Claude Code GELU task. |
+| `example_configs/quickstart_claude_mi325x.yaml` | Model-specific MI325X (`gfx942`) Claude Code GELU task. |
+| `example_configs/quickstart_claude_mi300a.yaml` | Model-specific MI300A (`gfx942`) Claude Code GELU task. |
 | `example_configs/quickstart_claude_mi355x.yaml` | One Claude Code GELU task on MI355X (`gfx950`); use this for a first run on MI355X. |
 | `example_configs/benchmark_cursor_mi355x.yaml` | Curated 60-task Cursor Agent benchmark on MI355X; use this for a longer benchmark only after installing and authenticating Cursor Agent. |
 
-Running `make docker-run` without `CONFIG` uses the MI300/MI300X Claude
-quickstart. On another GPU, pass the matching configuration explicitly; for
-example, use `CONFIG=example_configs/quickstart_claude_mi355x.yaml` on MI355X.
+Running `make docker-run` without `CONFIG` uses the generic MI300-series Claude
+quickstart. Pass the exact model configuration to give the optimization agent
+the correct XCD, CU, and HBM limits. For example, use
+`CONFIG=example_configs/quickstart_claude_mi325x.yaml` on MI325X.
 
 FlyDSL tasks require FlyDSL in the container. The pinned image may already provide it; otherwise run:
 

agents/mini_swe_triton/launch_agent.py

@@ -23,6 +23,7 @@
 import yaml
 
 from agents import register_agent
+from src.prompt_builder import _load_cheatsheet
 
 
 def _read_stream(stream, lines: list, prefix: str, log_func):
@@ -122,6 +123,16 @@ def launch_agent(eval_config: dict[str, Any], task_config_dir: str, workspace: s
     num_parallel = agent_config.get("agent", {}).get("num_parallel", 2)
     model = agent_config.get("agent", {}).get("model", "claude-opus-4-6")
     step_limit = agent_config.get("agent", {}).get("step_limit", 100)
+    target_gpu_model = str(eval_config.get("target_gpu_model") or "MI300")
+    project_root = Path(__file__).resolve().parents[2]
+    architecture_context, gfx_arch = _load_cheatsheet(
+        task_config.get("task_type", "triton2triton"),
+        target_gpu_model,
+        project_root,
+        task_config,
+        logger,
+        include_knowledge=False,
+    )
 
     # PYTHONPATH for mini-swe-agent modules. GEAK_SRC must point to the
     # absolute path of the GEAK source tree (the directory containing
@@ -149,6 +160,7 @@ def launch_agent(eval_config: dict[str, Any], task_config_dir: str, workspace: s
     logger.info(f"  num_parallel: {num_parallel}")
     logger.info(f"  model:        {model}")
     logger.info(f"  step_limit:   {step_limit}")
+    logger.info(f"  target_gpu:   {target_gpu_model} ({gfx_arch or 'unknown architecture'})")
     logger.info("=" * 60)
 
     all_output: list[str] = []
@@ -161,7 +173,11 @@ def launch_agent(eval_config: dict[str, Any], task_config_dir: str, workspace: s
     else:
         kernel_snippet = kernel_code
 
-    task_prompt = f"""Optimize this Triton GPU kernel for maximum performance on AMD MI300X (gfx942/gfx950).
+    target_description = target_gpu_model
+    if gfx_arch:
+        target_description += f" ({gfx_arch})"
+
+    task_prompt = f"""Optimize this Triton GPU kernel for maximum performance on AMD {target_description}.
 
 The kernel is at: {kernel_path.name}
 The test harness is at: {harness_path.name}
@@ -176,7 +192,9 @@ def launch_agent(eval_config: dict[str, Any], task_config_dir: str, workspace: s
 - Correctness must pass after your changes
 - Focus on real kernel-body optimizations (block sizes, memory access patterns,
   vectorization, loop unrolling, warp-level primitives)
-- Target: AMD MI300X with gfx942/gfx950 architecture, 304 CUs, HBM3
+
+Target hardware context:
+{architecture_context or f"No model profile is available for {target_description}; inspect the visible GPU before tuning."}
 
 Current kernel code:
 ```python

docs/examples/examples.md

@@ -17,11 +17,11 @@ multi-GPU examples use `make docker-parallel-run`.
 
 Run the single-task Claude Code quickstart for the physical GPU.
 
-1. Select the MI300/MI300X example (use
-   `example_configs/quickstart_claude_mi355x.yaml` on MI355X):
+1. Select the example matching the physical GPU (use the adjacent MI325X,
+   MI300A, or MI355X quickstart for those models):
 
     ```bash
-    CONFIG_PATH=example_configs/quickstart_claude_mi300.yaml
+    CONFIG_PATH=example_configs/quickstart_claude_mi300x.yaml
     ```
 
 2. Run:

docs/how-to/run-evaluation.md

@@ -15,18 +15,23 @@ resume, and inspect a run.
 ## Choose or create a run configuration
 
 A run configuration selects the agent, tasks, and target GPU. The repository
-ships three examples:
+ships generic and model-specific quickstarts plus a longer benchmark example:
 
 | Configuration | Purpose |
 | --- | --- |
-| `example_configs/quickstart_claude_mi300.yaml` | One Claude Code GELU task on MI300/MI300X (`gfx942`). |
+| `example_configs/quickstart_claude_mi300.yaml` | Generic MI300-series (`gfx942`) fallback. |
+| `example_configs/quickstart_claude_mi300x.yaml` | One Claude Code GELU task on MI300X (`gfx942`). |
+| `example_configs/quickstart_claude_mi325x.yaml` | One Claude Code GELU task on MI325X (`gfx942`). |
+| `example_configs/quickstart_claude_mi300a.yaml` | One Claude Code GELU task on MI300A (`gfx942`). |
 | `example_configs/quickstart_claude_mi355x.yaml` | One Claude Code GELU task on MI355X (`gfx950`). |
 | `example_configs/benchmark_cursor_mi355x.yaml` | Curated 60-task Cursor Agent benchmark on MI355X; use only after installing and authenticating Cursor Agent. |
 
 For a first run, select the quickstart that matches the physical GPU:
 
 ```bash
-CONFIG_PATH=example_configs/quickstart_claude_mi300.yaml
+CONFIG_PATH=example_configs/quickstart_claude_mi300x.yaml
+# MI325X: example_configs/quickstart_claude_mi325x.yaml
+# MI300A: example_configs/quickstart_claude_mi300a.yaml
 # For MI355X instead:
 # CONFIG_PATH=example_configs/quickstart_claude_mi355x.yaml
 ```

docs/install/install.md

@@ -93,24 +93,31 @@ their own runtime dependencies. Review the corresponding directory under
 ## Choose an example configuration
 
 Choose the configuration that matches the physical GPU and installed agent.
-The two quickstart configurations each run one GELU task; the benchmark
-configuration is a longer 60-task Cursor Agent run.
+Each quickstart configuration runs one GELU task; the benchmark configuration
+is a longer 60-task Cursor Agent run.
 
 | Configuration | Purpose |
 | --- | --- |
-| `example_configs/quickstart_claude_mi300.yaml` | First Claude Code run on MI300/MI300X (`gfx942`). |
+| `example_configs/quickstart_claude_mi300.yaml` | Generic MI300-series (`gfx942`) fallback. |
+| `example_configs/quickstart_claude_mi300x.yaml` | First Claude Code run on MI300X (`gfx942`). |
+| `example_configs/quickstart_claude_mi325x.yaml` | First Claude Code run on MI325X (`gfx942`). |
+| `example_configs/quickstart_claude_mi300a.yaml` | First Claude Code run on MI300A (`gfx942`). |
 | `example_configs/quickstart_claude_mi355x.yaml` | First Claude Code run on MI355X (`gfx950`). |
 | `example_configs/benchmark_cursor_mi355x.yaml` | Curated 60-task Cursor Agent benchmark on MI355X; requires an installed and authenticated Cursor Agent CLI. |
 
-The default `make docker-run` configuration is the MI300/MI300X quickstart.
-On another GPU, always pass the matching `CONFIG`; for MI355X, use
-`example_configs/quickstart_claude_mi355x.yaml`.
+The default `make docker-run` configuration is the generic MI300-series
+quickstart. Prefer the matching model-specific `CONFIG` so the optimization
+agent receives accurate hardware limits.
 
 Select one quickstart configuration in the current shell:
 
 ```bash
-# MI300/MI300X:
-CONFIG_PATH=example_configs/quickstart_claude_mi300.yaml
+# MI300X:
+CONFIG_PATH=example_configs/quickstart_claude_mi300x.yaml
+
+# For MI325X or MI300A instead:
+# CONFIG_PATH=example_configs/quickstart_claude_mi325x.yaml
+# CONFIG_PATH=example_configs/quickstart_claude_mi300a.yaml
 
 # For MI355X instead:
 # CONFIG_PATH=example_configs/quickstart_claude_mi355x.yaml

docs/reference/api-reference.md

@@ -19,7 +19,7 @@ A run configuration defines a single experiment. Start from a file under
 | --- | --- | --- |
 | `agent.template` | string | Agent to run. One of the [supported agents](../how-to/agents.md#supported-agents). |
 | `tasks` | list of strings | Task selectors relative to `tasks/`. Use `all` for every task, a category prefix for a group, or a full path for a single task. |
-| `target_gpu_model` | string | Target GPU model, for example `MI300` or `MI355X`. Used to select the Docker image architecture, set `PYTORCH_ROCM_ARCH`, and name the workspace. |
+| `target_gpu_model` | string | Target GPU model: `MI300X`, `MI325X`, `MI300A`, `MI355X`, or the generic `MI300` compatibility target. Used to select the Docker image architecture, compose model-aware optimization context, set `PYTORCH_ROCM_ARCH`, and name the workspace. |
 | `log_directory` | string | Directory for run logs. |
 | `workspace_directory_prefix` | string | Prefix for the workspace directory. The full name is `<prefix>_<gpu>_<agent>`. |
 
@@ -56,7 +56,7 @@ The in-container `main.py` entrypoint accepts these flags:
 
 | Flag | Description |
 | --- | --- |
-| `--config_name <file>` | Config file to load (default `example_configs/quickstart_claude_mi300.yaml` for MI300/MI300X). Pass a matching config explicitly on another GPU |
+| `--config_name <file>` | Config file to load (default `example_configs/quickstart_claude_mi300.yaml` for the generic MI300-series target). Pass a model-specific config for exact hardware context |
 | `--run-suffix <suffix>` | Suffix appended to the run directory name (letters, numbers, `.`, `_`, `-` only). Useful for labeling A/B runs |
 | `--resume-run <run_dir>` | Resume a specific run directory, skipping completed tasks |
 | `--resume-latest` | Resume the most recent run in the workspace |

docs/reference/compatibility-matrix.md

@@ -9,12 +9,23 @@ myst:
 
 ## Hardware
 
-The following hardware configurations are supported and tested.
+The following GPU models have explicit runtime and optimization-prompt
+configuration. Model-to-architecture and Docker routing are covered by automated
+tests. Hardware task validation is listed separately so that configured support
+is not mistaken for results from unavailable physical GPUs.
 
-| Component | Supported | Notes |
-| --- | --- | --- |
-| GPU architecture | AMD Instinct™ MI300 series | `target_gpu_model: MI300` |
-| GPU architecture | AMD Instinct™ MI355X | `target_gpu_model: MI355X` |
+| GPU model | Configuration value | LLVM target | Validation status |
+| --- | --- | --- | --- |
+| AMD Instinct™ MI300X | `target_gpu_model: MI300X` | `gfx942` | Supported by the existing MI300-series runtime path. |
+| AMD Instinct™ MI325X | `target_gpu_model: MI325X` | `gfx942` | Configuration tested; physical GPU task validation pending. |
+| AMD Instinct™ MI300A | `target_gpu_model: MI300A` | `gfx942` | Configuration tested; physical GPU task validation pending. |
+| Generic AMD Instinct™ MI300 series | `target_gpu_model: MI300` | `gfx942` | Backward-compatible fallback; runtime detection is required for SKU-specific limits. |
+| AMD Instinct™ MI355X | `target_gpu_model: MI355X` | `gfx950` | Supported and tested. |
+
+MI300X, MI325X, and MI300A share the `gfx942` compiler target but load separate
+hardware profiles. The shared CDNA 3 guidance contains ISA-level constraints;
+each profile supplies its model-specific compute topology, HBM limits, and
+partitioning guidance.
 
 ## Software
 

docs/reference/release-notes.md

@@ -27,8 +27,10 @@ experimentation platform and RL-ready task environment for GPU kernel agents.
 - **Docker-first execution**: select pinned ROCm/SGLang images from the target
   GPU architecture.
 - **Example run configurations**: provide one-task Claude Code quickstarts for
-  MI300/MI300X and MI355X plus a curated 60-task Cursor benchmark for MI355X
-  under `example_configs/`.
+  the generic MI300-series target, MI300X, MI325X, MI300A, and MI355X plus a
+  curated 60-task Cursor benchmark for MI355X under `example_configs/`.
+- **Model-aware MI300 profiles**: compose shared CDNA 3 (`gfx942`) guidance with
+  separate MI300X, MI325X, and MI300A topology and memory profiles.
 - **Centralized outcomes**: independently measure compilation, correctness, and
   GPU performance, then compute a configurable score.
 - **Visualization module**: build and serve the comparison dashboard through

example_configs/quickstart_claude_mi300.yaml

@@ -1,4 +1,5 @@
-# Minimal first-run example for Claude Code on MI300/MI300X (gfx942).
+# Generic MI300-series Claude Code quickstart (gfx942).
+# Prefer the model-specific MI300X, MI325X, or MI300A config for accurate tuning context.
 agent:
   template: claude_code
 

example_configs/quickstart_claude_mi300a.yaml

@@ -0,0 +1,10 @@
+# Minimal first-run example for Claude Code on MI300A (gfx942).
+agent:
+  template: claude_code
+
+tasks:
+  - hip2hip/gpumode/GELU
+
+target_gpu_model: MI300A
+log_directory: logs
+workspace_directory_prefix: workspace

example_configs/quickstart_claude_mi300x.yaml

@@ -0,0 +1,10 @@
+# Minimal first-run example for Claude Code on MI300X (gfx942).
+agent:
+  template: claude_code
+
+tasks:
+  - hip2hip/gpumode/GELU
+
+target_gpu_model: MI300X
+log_directory: logs
+workspace_directory_prefix: workspace

example_configs/quickstart_claude_mi325x.yaml

@@ -0,0 +1,10 @@
+# Minimal first-run example for Claude Code on MI325X (gfx942).
+agent:
+  template: claude_code
+
+tasks:
+  - hip2hip/gpumode/GELU
+
+target_gpu_model: MI325X
+log_directory: logs
+workspace_directory_prefix: workspace

main.py

@@ -38,8 +38,8 @@
     default="example_configs/quickstart_claude_mi300.yaml",
     help=(
         "run configuration for AgentKernelArena (default: "
-        "example_configs/quickstart_claude_mi300.yaml for MI300/MI300X). "
-        "Select a matching config explicitly when using another GPU."
+        "example_configs/quickstart_claude_mi300.yaml for the generic MI300-series target). "
+        "Select a model-specific config for exact hardware tuning context."
     ),
 )
 parser.add_argument(

src/prompt_builder.py

@@ -79,17 +79,19 @@ def task_result_format(task_result_template: str) -> str:
 
 
 def _load_cheatsheet(task_type_name: str, target_gpu_model: str, project_root: Path,
-                     task_config: dict, logger: logging.Logger) -> tuple[str, str | None]:
+                     task_config: dict, logger: logging.Logger, *,
+                     include_knowledge: bool = True) -> tuple[str, str | None]:
     """
     Load the combined cheatsheet prompt and resolve the target gfx arch string.
 
     The cheatsheet config (default_cheatsheet.yaml) has two independent sections:
-      - architecture: maps GPU model name → hardware spec document + gfx_arch string
+      - architecture: maps GPU model name → hardware spec document(s) + gfx_arch string
       - knowledge:    maps target language → language best-practice guide (GPU-agnostic)
 
     For task_type ``repository``, the task config must set ``repository_language`` to a
     key under ``knowledge`` (e.g. hip, triton). Add new stacks by extending that map and the
-    referenced markdown file.
+    referenced markdown file. Set ``include_knowledge=False`` for integrations that need
+    only the target architecture and model profile.
 
     Returns:
         (cheatsheet_text, gfx_arch)  where gfx_arch may be None if not found.
@@ -116,16 +118,27 @@ def _load_cheatsheet(task_type_name: str, target_gpu_model: str, project_root: P
         )
         if arch_entry:
             gfx_arch = arch_entry.get('gfx_arch')
-            arch_file = arch_entry.get('file')
-            if arch_file:
+            arch_files = arch_entry.get('files')
+            if arch_files is None:
+                legacy_arch_file = arch_entry.get('file')
+                arch_files = [legacy_arch_file] if legacy_arch_file else []
+            elif isinstance(arch_files, str):
+                arch_files = [arch_files]
+
+            for arch_file in arch_files:
                 arch_path = project_root / arch_file
                 parts.append(arch_path.read_text())
                 logger.info(f"Loaded architecture context for '{target_gpu_model}': {arch_path}")
-            else:
-                logger.warning(f"Architecture entry for '{target_gpu_model}' has no 'file' key")
+            if not arch_files:
+                logger.warning(
+                    f"Architecture entry for '{target_gpu_model}' has no 'file' or 'files' key"
+                )
         else:
             logger.warning(f"No architecture entry for GPU '{target_gpu_model}' in default_cheatsheet.yaml")
 
+        if not include_knowledge:
+            return "\n\n---\n\n".join(parts) if parts else "", gfx_arch
+
         # --- Knowledge section ---
         # L3 repository tasks: language is configured per task (repository_language → key in default_cheatsheet.yaml).
         if task_type_name == 'repository':