Skip to content

Update: convert paged attention example from float16 to bfloat16 #4

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
chenshengxin2026 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Update: convert paged attention example from float16 to bfloat16 #4

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
186fe51
Build both a2a3 and a5 onboard variants; ci.py auto-runs all sim plat…
hw-native-sys-bot Apr 8, 2026
9b9ed4d
Update: adapt simpler to pto-isa bfloat16 CPU-sim support
chenshengxin2026 Apr 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions .github/workflows/ci.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,7 @@ jobs:
- name: Set up C++ compiler
run: |
if [[ "${{ runner.os }}" == "Linux" ]]; then
sudo add-apt-repository -y ppa:ubuntu-toolchain-r/test
sudo apt-get update
sudo apt-get install -y ninja-build
sudo apt-get install -y g++-15 || sudo apt-get install -y g++
Expand Down Expand Up @@ -131,7 +132,7 @@ jobs:
run: pip install .

- name: Run simulation examples (a2a3sim)
run: python ci.py -p a2a3sim -c 882c4db -t 600 --clone-protocol https
run: python ci.py -p a2a3sim -c 8830244b -t 600 --clone-protocol https

st-sim-a5:
runs-on: ${{ matrix.os }}
Expand All @@ -148,6 +149,7 @@ jobs:
- name: Set up C++ compiler
run: |
if [[ "${{ runner.os }}" == "Linux" ]]; then
sudo add-apt-repository -y ppa:ubuntu-toolchain-r/test
sudo apt-get update
sudo apt-get install -y ninja-build
sudo apt-get install -y g++-15 || sudo apt-get install -y g++
Expand Down Expand Up @@ -181,7 +183,7 @@ jobs:
run: pip install .

- name: Run simulation examples (a5sim)
run: python ci.py -p a5sim -c 882c4db -t 600 --clone-protocol https
run: python ci.py -p a5sim -c 8830244b -t 600 --clone-protocol https

# ---------- Python unit tests (a2a3 hardware) ----------
ut-py-a2a3:
Expand Down Expand Up @@ -215,7 +217,7 @@ jobs:
- name: Run on-device examples (a2a3)
run: |
export PATH="$HOME/.local/bin:$PATH"
source ${ASCEND_HOME_PATH}/bin/setenv.bash && python ci.py -p a2a3 -d ${DEVICE_RANGE} -c 882c4db -t 600 --clone-protocol https
source ${ASCEND_HOME_PATH}/bin/setenv.bash && python ci.py -p a2a3 -d ${DEVICE_RANGE} -c 8830244b -t 600 --clone-protocol https


# ---------- Detect A5 changes (runs on GitHub server, not A5 machine) ----------
Expand Down Expand Up @@ -283,11 +285,11 @@ jobs:
- name: Build nanobind extension
run: |
source ${ASCEND_HOME_PATH}/bin/setenv.bash
PTO_ARCH=a5 pip install .
pip install .

- name: Run on-device examples (a5)
run: |
export PATH="$HOME/.local/bin:$PATH"
source ${ASCEND_HOME_PATH}/bin/setenv.bash
DEVICE_LIST=$(python -c "s,e='${DEVICE_RANGE}'.split('-'); print(','.join(str(i) for i in range(int(s),int(e)+1)))")
task-submit --timeout 1800 --max-time 1800 --device "$DEVICE_LIST" --run "python ci.py -p a5 -d ${DEVICE_RANGE} -c 882c4db -t 1200 --clone-protocol https"
task-submit --timeout 1800 --max-time 1800 --device "$DEVICE_LIST" --run "python ci.py -p a5 -d ${DEVICE_RANGE} -c 8830244b -t 1200 --clone-protocol https"
103 changes: 67 additions & 36 deletions ci.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,9 @@
per device, reusing ChipWorker across tasks that share the same runtime.

Usage:
python tools/ci.py -p a2a3 -d 5-8 -c 6622890 -t 600
python tools/ci.py -p a2a3sim -r tensormap_and_ringbuffer -c 6622890 -t 600
python ci.py # all sim platforms
python ci.py -p a2a3sim -r tensormap_and_ringbuffer -c 6622890 # single platform
python ci.py -p a2a3 -d 5-8 -c 6622890 -t 600 # hardware with devices
"""

from __future__ import annotations
Expand Down Expand Up @@ -1100,7 +1101,7 @@ def _discover_valid_platforms() -> list[str]:

def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Batch CI test runner with ChipWorker reuse")
parser.add_argument("-p", "--platform", required=True)
parser.add_argument("-p", "--platform", default=None)
parser.add_argument("-d", "--device", dest="device_range", default="0")
parser.add_argument("-r", "--runtime", default=None)
parser.add_argument(
Expand Down Expand Up @@ -1149,44 +1150,34 @@ def _watchdog_handler(signum, frame):
signal.signal(signal.SIGALRM, previous_handler)


def main() -> int:
logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s", force=True)

args = parse_args()
args.devices = parse_device_range(args.device_range)
def _run_single_platform(platform: str, args: argparse.Namespace) -> list[TaskResult]:
"""Run all tasks for a single platform. Returns list of TaskResults."""
is_sim = platform.endswith("sim")

valid_platforms = _discover_valid_platforms()
if valid_platforms and args.platform not in valid_platforms:
print(f"Unknown platform: {args.platform}")
print(f"Valid platforms: {' '.join(valid_platforms)}")
return 1

is_sim = args.platform.endswith("sim")

# Device-worker sub-command
if args.device_worker:
return device_worker_main(args)

# Step 1: Discover tasks
tasks = discover_tasks(args.platform, runtime_filter=args.runtime)
tasks = discover_tasks(platform, runtime_filter=args.runtime)
if not tasks:
logger.info("No tasks found")
return 0
logger.info(f"Discovered {len(tasks)} tasks")
logger.info(f"[{platform}] No tasks found")
return []
logger.info(f"[{platform}] Discovered {len(tasks)} tasks")

# Step 2: Compile and run.
# Compile and run.
# Both sim and hw use subprocess isolation (different runtimes cannot share a process).
# Within each subprocess, tasks with the same runtime share a ChipWorker.
# Override platform in args for subprocess spawning.
sub_args = argparse.Namespace(**vars(args))
sub_args.platform = platform
if is_sim:
all_results = _run_with_timeout("initial pass", args.timeout, lambda: run_sim_tasks_subprocess(tasks, args))
all_results = _run_with_timeout(
f"{platform} initial pass", args.timeout, lambda: run_sim_tasks_subprocess(tasks, sub_args)
)
else:
all_results = _run_with_timeout(
"initial pass",
f"{platform} initial pass",
args.timeout,
lambda: run_hw_tasks_subprocess(tasks, args.devices, args),
lambda: run_hw_tasks_subprocess(tasks, args.devices, sub_args),
)

# Step 3: Pin retry — re-run failed tasks with pinned PTO-ISA commit.
# Pin retry — re-run failed tasks with pinned PTO-ISA commit.
final: dict[str, TaskResult] = {}
for r in all_results:
final[r.name] = r
Expand All @@ -1195,27 +1186,67 @@ def main() -> int:
if failures and args.pto_isa_commit:
failed_names = {r.name for r in failures}
failed_tasks = [t for t in tasks if t.name in failed_names]
logger.info(f"[CI] {len(failed_tasks)} failure(s), retrying with pinned PTO-ISA {args.pto_isa_commit}")
logger.info(f"[{platform}] {len(failed_tasks)} failure(s), retrying with pinned PTO-ISA {args.pto_isa_commit}")
if is_sim:
pin_results = _run_with_timeout(
"pin retry",
f"{platform} pin retry",
args.timeout,
lambda: run_sim_tasks_subprocess(failed_tasks, args, pto_isa_commit=args.pto_isa_commit),
lambda: run_sim_tasks_subprocess(failed_tasks, sub_args, pto_isa_commit=args.pto_isa_commit),
)
else:
pin_results = _run_with_timeout(
"pin retry",
f"{platform} pin retry",
args.timeout,
lambda: run_hw_tasks_subprocess(
failed_tasks,
args.devices,
args,
sub_args,
pto_isa_commit=args.pto_isa_commit,
),
)
all_results.extend(pin_results)

# Step 4: Summary
return all_results


def main() -> int:
logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s", force=True)

args = parse_args()
args.devices = parse_device_range(args.device_range)

valid_platforms = _discover_valid_platforms()

# Device-worker sub-command (always needs explicit -p)
if args.device_worker:
if not args.platform:
print("--device-worker requires -p/--platform")
return 1
return device_worker_main(args)

# Determine which platforms to run
if args.platform:
if args.platform not in valid_platforms:
print(f"Unknown platform: {args.platform}")
print(f"Valid platforms: {' '.join(valid_platforms)}")
return 1
platforms = [args.platform]
else:
# No -p: run all sim platforms
platforms = [p for p in valid_platforms if p.endswith("sim")]
if not platforms:
print("No sim platforms found")
return 1
logger.info(f"No platform specified, running all sim platforms: {', '.join(platforms)}")

all_results: list[TaskResult] = []
for platform in platforms:
all_results.extend(_run_single_platform(platform, args))

if not all_results:
logger.info("No tasks found")
return 0

return print_summary(all_results)


Expand Down
76 changes: 38 additions & 38 deletions docs/ci.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,14 +8,14 @@ Design principles:

1. **Separate jobs per test category** — st, ut-py, and ut-cpp run as independent jobs for parallelism and clear dashboard visibility.
2. **Runner matches hardware tier** — no-hardware tests run on `ubuntu-latest`; platform-specific tests run on self-hosted runners with the matching label (`a2a3`, `a5`).
3. **sim vs onboard for st** — scene tests split into sim jobs (github-hosted, `ci.sh -p a2a3sim`) and onboard jobs (self-hosted, `ci.sh -p a2a3`).
3. **sim vs onboard for st** — scene tests split into sim jobs (github-hosted, `ci.py -p a2a3sim`) and onboard jobs (self-hosted, `ci.py -p a2a3`).

## Full Job Matrix

The complete test-type × hardware-tier matrix. Empty cells have no tests yet; only non-empty jobs exist in `ci.yml`.

| | github-hosted | a2a3 runner | a5 runner |
|---|---|---|---|
| Category | github-hosted | a2a3 runner | a5 runner |
| -------- | ------------- | ----------- | --------- |
| **ut-py** | `ut-py` | `ut-py-a2a3` | `ut-py-a5` |
| **ut-cpp** | *(empty)* | *(empty)* | *(empty)* |
| **st** | `st-sim-a2a3`, `st-sim-a5` | `st-onboard-a2a3` | `st-onboard-a5` |
Expand All @@ -24,7 +24,7 @@ The complete test-type × hardware-tier matrix. Empty cells have no tests yet; o

Currently active jobs (a5 jobs commented out — no runner yet):

```
```text
PullRequest
├── ut-py (ubuntu-latest)
├── st-sim-a2a3 (ubuntu + macOS)
Expand All @@ -36,14 +36,14 @@ PullRequest
```

| Job | Runner | What it runs |
|-----|--------|-------------|
| --- | ------ | ------------ |
| `ut-py` | `ubuntu-latest` | `pytest tests -m "not requires_hardware" -v` |
| `st-sim-a2a3` | `ubuntu-latest`, `macos-latest` | `ci.sh -p a2a3sim` |
| `st-sim-a5` | `ubuntu-latest`, `macos-latest` | `ci.sh -p a5sim` |
| `st-sim-a2a3` | `ubuntu-latest`, `macos-latest` | `ci.py -p a2a3sim` |
| `st-sim-a5` | `ubuntu-latest`, `macos-latest` | `ci.py -p a5sim` |
| `ut-py-a2a3` | a2a3 self-hosted | `pytest tests -m requires_hardware --platform a2a3 -v` |
| `st-onboard-a2a3` | a2a3 self-hosted | `ci.sh -p a2a3 -d ... --parallel` |
| `st-onboard-a2a3` | a2a3 self-hosted | `ci.py -p a2a3 -d ...` |
| `ut-py-a5` | a5 self-hosted | `pytest tests -m requires_hardware --platform a5 -v` |
| `st-onboard-a5` | a5 self-hosted | `ci.sh -p a5 -d ... --parallel` |
| `st-onboard-a5` | a5 self-hosted | `ci.py -p a5 -d ...` |

### Scheduling constraints

Expand All @@ -56,7 +56,7 @@ PullRequest
Three hardware tiers, applied to all test categories. See [testing.md](testing.md#hardware-classification) for the full table including per-category mechanisms (pytest markers, ctest labels, folder structure).

| Tier | CI Runner | Job examples |
|------|-----------|-------------|
| ---- | --------- | ------------ |
| No hardware | `ubuntu-latest` | `ut-py`, `st-sim-*` |
| Platform-specific (a2a3) | `[self-hosted, a2a3]` | `ut-py-a2a3`, `st-onboard-a2a3` |
| Platform-specific (a5) | `[self-hosted, a5]` | `ut-py-a5`, `st-onboard-a5` |
Expand All @@ -68,24 +68,24 @@ Three hardware tiers, applied to all test categories. See [testing.md](testing.m
Python unit tests. Run via pytest with marker filtering.

| File | Content | Hardware? |
|------|---------|-----------|
| ---- | ------- | --------- |
| `test_task_interface.py` | nanobind extension API tests | No |
| `test_runtime_builder.py` (mocked classes) | RuntimeBuilder discovery, error handling, build logic | No |
| `test_runtime_builder.py::TestRuntimeBuilderIntegration` | Real compilation across platform × runtime | Yes (`@pytest.mark.requires_hardware`) |

### `examples/` — Small examples (sim + onboard)

Small, fast examples that run on both simulation and real hardware. Organized as `examples/{arch}/{runtime}/{name}/`. Discovered and executed by `ci.sh`.
Small, fast examples that run on both simulation and real hardware. Organized as `examples/{arch}/{runtime}/{name}/`. Discovered and executed by `ci.py`.

### `tests/st/` — Scene tests (onboard only)

Large-scale, feature-rich hardware tests. Too slow or using instructions unsupported by the simulator. Organized as `tests/st/{arch}/{runtime}/{name}/`. Discovered and executed by `ci.sh` only when the platform is not sim.
Large-scale, feature-rich hardware tests. Too slow or using instructions unsupported by the simulator. Organized as `tests/st/{arch}/{runtime}/{name}/`. Discovered and executed by `ci.py` only when the platform is not sim.

### Shared structure

Both `examples/` and `tests/st/` cases follow the same layout:

```
```text
{name}/
golden.py # generate_inputs() + compute_golden()
kernels/
Expand Down Expand Up @@ -143,38 +143,39 @@ python tools/test_catalog.py cases --platform a2a3sim --source example
python tools/test_catalog.py cases --platform a2a3 --source st --format json
```

## `ci.sh` — Scene Test Runner
## `ci.py` — Scene Test Runner

`ci.sh` handles scene test execution (examples + st). It does **not** run pytest tests — those are invoked directly by the CI workflow.
`ci.py` handles scene test execution (examples + st). It does **not** run pytest tests — those are invoked directly by the CI workflow.

### Key features

- **Parallel execution**: `--parallel` runs sim tasks concurrently; hardware tasks use a shared device queue with `flock`-based locking.
- **Device queue**: Hardware tasks are distributed across devices specified by `-d`. Workers pop tasks from a shared queue atomically.
- **Retry**: Failed tasks are retried up to 3 times. Hardware workers quarantine a device after exhausting retries.
- **PTO-ISA pinning**: `-c <commit>` pins the PTO-ISA dependency. On first failure, ci.sh cleans the cached clone and retries with the pinned commit.
- **ChipWorker reuse**: Tasks sharing the same runtime reuse a single ChipWorker within their subprocess, avoiding repeated device init/teardown.
- **Subprocess isolation**: Different runtimes run in separate subprocesses (the host `.so` cannot be unloaded within a single process).
- **Device queue**: Hardware tasks are distributed across devices specified by `-d`. Workers pop tasks from a shared queue via threads.
- **Retry**: Failed tasks are retried up to 3 times. Hardware workers quarantine a device after a failure.
- **PTO-ISA pinning**: `-c <commit>` pins the PTO-ISA dependency. On first failure, re-runs failed tasks with the pinned commit.
- **Watchdog**: `-t <seconds>` sets a timeout. The entire run is aborted if it exceeds the limit.
- **Summary table**: After all tasks complete, a formatted results table is printed with pass/fail status, timing, device, and attempt count.

### Usage

```bash
# Simulation (github-hosted)
./ci.sh -p a2a3sim -t 600
# All sim platforms (no -p: auto-discovers a2a3sim, a5sim, etc.)
python ci.py -t 600

# Single sim platform
python ci.py -p a2a3sim -c 6622890 -t 600

# Hardware with device range
./ci.sh -p a2a3 -d 4-7 --parallel -t 600
python ci.py -p a2a3 -d 4-7 -c 6622890 -t 600

# Filter by runtime
./ci.sh -p a2a3sim -r tensormap_and_ringbuffer

# Pin PTO-ISA commit
./ci.sh -p a2a3sim -c 6622890
python ci.py -p a2a3sim -r tensormap_and_ringbuffer
```

### Task discovery

`ci.sh` scans two directories:
`ci.py` scans two directories:

1. `examples/` — included for both sim and onboard platforms.
2. `tests/st/` — included only for onboard platforms (non-sim).
Expand All @@ -183,14 +184,13 @@ For each directory, it walks subdirectories looking for `kernels/kernel_config.p

### Execution flow

```
1. Parse arguments (-p, -d, --parallel, -r, -c, -t)
2. Discover platforms and runtimes from src/
3. Discover tasks from examples/ and tests/st/
4. Run sim tasks (parallel or sequential)
└── On failure + -c flag: pin PTO-ISA, retry
5. Run hardware tasks (device queue with workers)
└── On failure + -c flag: pin PTO-ISA, retry
6. Print summary table
7. Exit 0 if all passed, 1 otherwise
```text
1. Parse arguments (-p, -d, -r, -c, -t)
2. If no -p: auto-discover all sim platforms and run each
3. For each platform:
a. Discover tasks from examples/ and tests/st/
b. Run tasks (subprocess per runtime group for sim, device queue for hw)
└── On failure + -c flag: pin PTO-ISA, retry failed tasks
4. Print combined summary table
5. Exit 0 if all passed, 1 otherwise
```
7 changes: 6 additions & 1 deletion docs/developer-guide.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,11 @@ Runtime binaries (host `.so`, aicpu `.so`, aicore `.o`) are pre-built during `pi

Persistent cmake build directories under `build/cache/` enable incremental compilation — only changed files are recompiled.

**Architecture note:** a2a3 and a5 differ only at runtime (device selection, block dimensions, etc.). The compiled binaries are architecture-independent — the same toolchain and flags produce artifacts that work on both chips. Therefore `pip install .` should build **all** architectures (both a2a3 and a5, both onboard and sim) whenever the corresponding toolchain is available. Toolchain detection (`build_runtimes.py`):

- **sim** (a2a3sim, a5sim): requires `gcc` + `g++` in `PATH`
- **onboard** (a2a3, a5): requires `ccec` in `PATH` + cross-compiler under `ASCEND_HOME_PATH`

### User code (per-example)

1. `python/kernel_compiler.py` — compiles user-written kernel `.cpp` files (one per `func_id`)
Expand Down Expand Up @@ -123,7 +128,7 @@ Run with: `python examples/scripts/run_example.py -k <kernels_dir> -g <golden.py
pip install -e .
```

This builds the nanobind `_task_interface` extension **and** pre-builds all runtime binaries for available toolchains into `build/lib/`. On x86_64, this means sim platforms only; on aarch64 hardware, onboard variants are also built.
This builds the nanobind `_task_interface` extension **and** pre-builds all runtime binaries for available toolchains into `build/lib/`. Sim platforms (a2a3sim, a5sim) are built when `gcc`/`g++` are available; onboard platforms (a2a3, a5) are built when `ccec` and the cross-compiler under `ASCEND_HOME_PATH` are available. Since a2a3 and a5 share the same compilation — differing only at runtime — both architectures are always built together when their toolchain is present.

### When to rebuild

Expand Down
Loading
Loading