Skip to content

Repository files navigation

CHASE

CHASE: Cross-Layer Heterogeneous Agile System-Exploration for Application-Driven Architecture is a research prototype for exploring physically feasible heterogeneous AI/HPC systems. It models candidate systems as hierarchical hardware graphs, maps application DAGs onto each candidate, evaluates the mapped traces with a calibration-aware event-driven simulator, and uses telemetry-guided search to propose improved architectures.

The repository provides a workload mapper, an ASTRA-sim-based congestion-aware simulator, a topology optimizer, calibration utilities, and end-to-end tools that connect them into a runnable mapper -> simulator -> optimizer workflow.

Capabilities

CHASE targets Cross-layer Heterogeneous Systems (XHS): systems that combine GPUs, CPUs, memory tiers, switches, and package/node/rack/cluster fabrics under power, budget, rack-space, and wiring constraints.

The tool suite provides:

  • Hierarchical hardware graphs represent packages, nodes, racks, clusters, typed components, explicit links, rank placement, and remote memory providers.
  • Constraint filtering rejects invalid candidates before expensive mapping and simulation.
  • Decoupled two-level exploration resolves the topology-mapping dependency: the mapper builds a topology-aware event trace for each candidate, while the optimizer evolves the hardware design using simulator telemetry.
  • Measurement-driven calibration anchors compute and communication timing to real platforms before large-scale projection.
  • Workload suites and cross-evaluation compare workload-specific and shared topology designs across sparse and LLM inputs.

Repository Layout

.
├── mapper/       C++ workload mapper; workload DAG + hardware projection -> taskflow / Chakra ET
├── simulator/    ASTRA-sim extension with explicit hardware graphs and remote memory telemetry
├── optimizer/    Python topology optimizer: search, exhaustive, TCRO, and TG-RL
├── tools/        End-to-end wrapper, topology examples, visualization, and scoring utilities
├── calibration/  Measured platform profiles and fitted calibration model artifacts
├── scripts/      Calibration, replay, validation, and cluster helper scripts
└── docs/         Root user manual for building, running, and understanding CHASE

mapper, simulator, and optimizer are git submodules. Initialize them before building.

Architecture Overview

workload DAG or LLM config
        |
        v
v2 hardware-topology candidate
        |
        v
tools/run_mapper_sim_pipeline.py
        |
        +--> mapper hardware projection
        +--> mapper taskflow
        +--> per-rank Chakra ET
        +--> AstraSim_Analytical_Congestion_Aware
        |
        v
run_summary.json + simulator stdout telemetry
        |
        v
optimizer feedback and next topology candidate

The unified topology is the hardware source of truth. Mapper sees only schedulable rank compute nodes. The simulator sees the full explicit graph, including switches, memory-only nodes, directed links, routing paths, queueing, and remote memory service. The optimizer normally exports candidates identified as terrapod.hardware_topology.v2 and evaluates them through the same wrapper.

Quick Start

The commands below run a minimal workflow for validating the installed toolchain.

1. Fetch Submodules

git submodule update --init --recursive

2. Install System Dependencies

On Ubuntu/Debian:

sudo apt-get update
sudo apt-get install -y \
  build-essential \
  cmake \
  protobuf-compiler \
  libprotobuf-dev \
  libabsl-dev \
  python3 \
  python3-venv \
  python3-pip \
  graphviz

On macOS with Homebrew:

brew install cmake protobuf abseil python@3.12 graphviz

3. Build Mapper

cd mapper
make
cd ..

test -x mapper/mapper_demo

4. Build Simulator

cd simulator
build/astra_analytical/build.sh
cd ..

test -x simulator/build/astra_analytical/build/bin/AstraSim_Analytical_Congestion_Aware

5. Install Optimizer

cd optimizer
python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip setuptools wheel
pip install -e ".[dev]"
cd ..

For TG-RL v2, install the RL extra:

cd optimizer
source .venv/bin/activate
pip install -e ".[dev,rl]"
cd ..

On CPU-only machines, install the CPU PyTorch wheel first to avoid pulling a large CUDA-enabled wheel:

cd optimizer
source .venv/bin/activate
pip install "torch>=2.3" --index-url https://download.pytorch.org/whl/cpu
pip install -e ".[dev,rl]"
cd ..

6. Run a Minimal Mapper -> Simulator Validation

python3 tools/run_mapper_sim_pipeline.py \
  --topology tools/examples/topologies/tc_2gpu_1cpu_switch_memory_pool.json \
  --topology-format hardware \
  --workload mapper/examples/cg_iteration_workload.json \
  --mapper heft \
  --parallel none \
  --save-wrapper-inputs \
  --out /tmp/chase_wrapper_validation

Check:

/tmp/chase_wrapper_validation/outputs/run_summary.json
/tmp/chase_wrapper_validation/outputs/mapper_stdout.txt
/tmp/chase_wrapper_validation/outputs/simulator_stdout.txt

7. Run an Optional Minimal Calibrated Validation

The repository includes calibration profiles and a fitted continuous model. Pass the model through the wrapper when you want mapper and simulator inputs materialized from calibrated compute and communication scales.

python3 tools/run_mapper_sim_pipeline.py \
  --topology tools/examples/topologies/tc_2gpu_1cpu_switch_memory_pool.json \
  --topology-format hardware \
  --workload mapper/examples/cg_iteration_workload.json \
  --mapper heft \
  --parallel none \
  --calibration-fit-model calibration/calibration_fit_model.json \
  --save-wrapper-inputs \
  --out /tmp/chase_calibrated_validation

For collecting new measurements, fitting the continuous model, and validating holdout platforms, see Calibration.

8. Run a Minimal Optimizer Validation

source optimizer/.venv/bin/activate

codesign-opt search \
  --catalog optimizer/examples/component_catalog.json \
  --space optimizer/examples/search_space.json \
  --workload mapper/examples/cg_iteration_workload.json \
  --generations 1 \
  --population 2 \
  --concurrency 1 \
  --out /tmp/chase_search_validation

Optional minimal TG-RL v2 validation after installing .[dev,rl]:

codesign-opt tgrl \
  --catalog optimizer/examples/component_catalog_tgrl.json \
  --space optimizer/examples/search_space_tgrl.json \
  --workload mapper/examples/cg_iteration_workload.json \
  --episodes 1 \
  --steps-per-episode 1 \
  --mode v2 \
  --concurrency 1 \
  --ppo-epochs 1 \
  --device auto \
  --out /tmp/chase_tgrl_v2_validation

Documentation

Start with the root manual:

Submodule-level documents remain useful as deep references, especially:

License and Attribution

Original CHASE software and documentation are Copyright (c) 2026 Yuchen Fan, Minghong Sun, Jikui Ma, and Shunyu Mao, and are released under the MIT License. See Authors and Copyright Holders for the project attribution.

The simulator is derived from ASTRA-sim 2.0. ASTRA-sim and all other third-party components retain their original copyright and license terms. See Third-Party Notices before redistributing the repository or simulator.

About

superpod

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages