Add: L3 communication-computation example - #869
Conversation
📝 WalkthroughWalkthroughThe PR adds an N-rank L3 all-gather plus GEMM example. It implements tiled communication and GEMM kernels, rank orchestration, CLI validation, correctness checks, benchmarking, and documentation updates. ChangesN-rank L3 all-gather + GEMM
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant HostOrchestration
participant ChipOrchestration
participant CommKernel
participant GemmKernel
CLI->>HostOrchestration: parse execution options
HostOrchestration->>ChipOrchestration: allocate per-rank buffers and launch rank
ChipOrchestration->>CommKernel: publish local shards and signals
ChipOrchestration->>GemmKernel: wait for gathered chunks and compute tiles
GemmKernel-->>CLI: return output for correctness and benchmark reporting
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Add the new N-rank L3 all-gather + GEMM example to the advanced and index example pages, add the `# ci: devices=2` marker it needs to run on the real a2a3 CI job (without it, task-submit borrows only one card and the script's own world_size>=2 check fails), and align its "chunked" wording with the "tiled" terminology used elsewhere.
Shenggan
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/advanced/l3_allgather_gemm.py`:
- Around line 234-240: Reset the persistent signal buffer before each repeated
benchmark launch in the kernel containing the all-gather loop and signal window
setup. Ensure each round clears or reinitializes signal values before the
`pld.system.wait` checks expecting 1, without changing the existing gathered
buffer or rank orchestration behavior.
- Around line 207-225: Update argument parsing for the comm_cores and gemm_cores
options to validate that their sum does not exceed the per-chip platform core
budget, in addition to the existing task-level checks. Reject invalid values
before overlap submissions run, using the platform’s established core-budget
symbol and preserving current validation for each option.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 78f018e5-fb6c-4b50-957d-088bb4fd8267
📒 Files selected for processing (3)
docs/examples/advanced.mddocs/examples/index.mdexamples/advanced/l3_allgather_gemm.py
| with pl.manual_scope(): | ||
| _gathered, _ = pl.spmd_submit( | ||
| self.comm_local_shard, | ||
| local_a, | ||
| gathered, | ||
| signal, | ||
| core_num=comm_cores, | ||
| ) | ||
| out, _ = pl.spmd_submit( | ||
| self.gemm_from_gathered, | ||
| local_a, | ||
| gathered, | ||
| signal, | ||
| weight, | ||
| out, | ||
| core_num=gemm_cores, | ||
| ) | ||
|
|
||
| return out |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check spmd_submit semantics inside manual_scope and core-count constraints.
set -euo pipefail
rg -n -C10 'def spmd_submit' --glob '*.py' || true
rg -n -C6 'def manual_scope' --glob '*.py' || true
# Other examples that issue two spmd_submit calls in one manual_scope.
rg -n -C12 'spmd_submit' examples || true
# Core budget validation elsewhere in the repo.
rg -n -C4 'core_num|core_budget|MAX_CORE|aic_num|aiv_num' --glob '*.py' || trueRepository: hw-native-sys/pypto-lib
Length of output: 161
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== files named pypto/CLAUDE.md and rules =="
fd -a 'CLAUDE\.md$|.*\.md$' .claude 2>/dev/null | sed 's#^\./##' | head -200 || true
wc -l .claude/CLAUDE.md 2>/dev/null || true
echo "== relevant pypto/coding style files =="
fd -a 'pypto-coding-style\.md$|.*\.md$' docs/pypto-coding 2>/dev/null | sed ' s#^\./##' || true
echo "== candidate files containing pypto imports/submit/manual_scope/gathered/core_num =="
rg -n -C5 'pl\.(manual_scope|spmd_submit)|spmd_submit|manual_scope|gathered|core_num' --glob '*.py' . || true
echo "== target file outline and relevant section =="
wc -l examples/advanced/l3_allgather_gemm.py 2>/dev/null || true
sed -n '1,280p' examples/advanced/l3_allgather_gemm.py 2>/dev/null || true
echo "== imports and namespace aliases in py files =="
rg -n 'import .* as pl|from .* import .* as pl|import pypto' --glob '*.py' . || trueRepository: hw-native-sys/pypto-lib
Length of output: 50380
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== target file head =="
sed -n '1,80p' examples/advanced/l3_allgather_gemm.py 2>/dev/null || true
echo "== target function argument parsing and chip_orch call =="
sed -n '266,380p' examples/advanced/l3_allgather_gemm.py 2>/dev/null || true
echo "== focused search for spmd_submit/manual_scope definitions =="
rg -n -C12 '^\s*def spmd_submit|^\s*async def spmd_submit|^\s*def manual_scope|^\s*async def manual_scope|^spmd_submit\s*=|manual_scope\s*=' \
--glob '*.py' . --glob '!examples/advanced/l3_allgather_gemm.py' --glob '!models/deepseek_v4_flash_mtp/*' --generic -g '!**/*.py' \
| head -300 || true
echo "== focused search for core budget validation =="
rg -n -C5 'aic_num|aiv_num|MAX_CORE|core_num|core_budget|num_cores|total.*core|device.*core' \
--glob '*.py' . || true
echo "== focused search for manual_scope usage around spmd =="
rg -n -C6 'pl\.manual_scope|with pl\.manual_scope|manual_scope' \
--glob '*.py' . | awk 'f<60{print; f++}' || trueRepository: hw-native-sys/pypto-lib
Length of output: 20565
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
path = Path("examples/advanced/l3_allgather_gemm.py")
text = path.read_text()
M_TILES = 2048 // 128
print(f"M_TILES = {M_TILES}")
print(f"current_default_core_sum = {4 + 8}")
platforms = ["a2a3", "a2a3sim", "a5", "a5sim"]
for platform in platforms:
# No platform-specific constants are referenced here.
if 4 + 8 <= M_TILES:
print(f"{platform}: default 4+8 <= M_TILES={M_TILES}")
required = "args.comm_cores > M_TILES"
print(f"comm_over_tiles_warns_condition_present = {required in text}")
required = "args.gemm_cores > gemm_tasks"
print(f("gemm_over_tasks_warns_condition_present = {required in text}"))
# Inspect dependency-like args around spmd_submit in this file.
lines = text.splitlines()
for i, line in enumerate(lines):
if "spmd_submit" in line:
print(f"spmd_submit at line {i+1}: {line.strip()}")
# Count manual_scope and spmd_submit occurrences in this file only.
print(f"manual_scope_count = {text.count('pl.manual_scope')}")
print(f"spmd_submit_count = {text.count('pl.spmd_submit')}")
PYRepository: hw-native-sys/pypto-lib
Length of output: 490
🌐 Web query:
pypto platform core budget a2a3 a5 aic_num aiv_num manual_scope spmd_submit
💡 Result:
PyPTO (Parallel Tensor/Tile Operation) is a high-performance programming framework for AI accelerators [1] that utilizes a tile-based programming model [1] and MPMD (Multiple Program Multiple Data) execution scheduling [1] to run task graphs on platforms such as Ascend A2/A3 (a2a3) and A5 (a5) [2][3]. The terms provided relate to specific runtime, architectural, and compilation configurations within the PyPTO ecosystem: - Platform Variants: a2a3 and a5 refer to distinct Ascend hardware targets (e.g., Ascend 910B and 950 respectively), which differ in architectural constants, memory reservation requirements, and ISA implementations [4][2][5][6]. - Core Budget: Generally refers to the allocation of AICore (AIC) and Vector Core (AIV) resources managed by the runtime for task scheduling [4][7][1]. - aic_num and aiv_num: These are configuration parameters or metrics defining the number of AICore and AIV units available or utilized per task or die, often used to determine scheduling and communication strategies [4]. - manual_scope: A programming mode in PyPTO where resources (such as tiles or buffers) must be explicitly bound or managed by the developer [6], often used to optimize performance or handle complex dependencies that the automatic runtime cannot infer [8][6]. - spmd_submit: Relates to the submission of SPMD (Single Program Multiple Data) style tasks within the framework, often involving the orchestration of fused kernels [9]. Explicit dependency management via submission APIs (e.g., pl.submit(..., deps=[...])) is frequently used as a manual override to resolve data hazards (such as WAR/RAW) that the automatic dependency tracker misses [8]. These components are part of the PTO2 (Parallel Task Orchestration v2) runtime system [7], which coordinates task graphs through host-side compilation, device-side AICPU-based scheduling, and execution on AICores [7][1]. Performance debugging within these environments often involves monitoring task flow, buffer reservations, and synchronization fences (e.g., MTE3→V pipes) to avoid deadlocks or suboptimal execution [8][9][7].
Citations:
- 1: https://github.com/hw-native-sys/pypto/
- 2: https://github.com/hw-native-sys/simpler
- 3: https://github.com/uv-xiao/pto-runtime
- 4: https://pto-isa.gitcode.com/docs/reference/pto-cvid-cluster-id-mapping/
- 5: https://gitcode.com/cann/pto-isa/tree/master/tests
- 6: https://pto-isa.gitcode.com/docs/isa/TPRELU/
- 7: https://github.com/xwhu/pypto-lib/blob/main/pto2_rt.md
- 8: [Bug] decode_attention_swa/hca need a no-op self-copy hack to force a WAR edge auto-dep misses — replace with pl.submit(deps=) manual dependency #481
- 9: [Bug] Removing a vector op (pl.transpose) from a fused cube+vec spmd reallocates the MTE3→V fence EVENT_ID and deadlocks at runtime (AICPU 507018) pypto#1789
Validate the per-rank core budget before running overlap submissions.
The example only checks task-level overuse: --comm-cores > M_TILES and --gemm-cores > gemm_tasks. It does not validate comm_cores + gemm_cores <= the per-chip platform core budget, so users can pass arbitrary positive values that exceed available cores. Add that validation or error at argument parsing time.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/advanced/l3_allgather_gemm.py` around lines 207 - 225, Update
argument parsing for the comm_cores and gemm_cores options to validate that
their sum does not exceed the per-chip platform core budget, in addition to the
existing task-level checks. Reject invalid values before overlap submissions
run, using the platform’s established core-budget symbol and preserving current
validation for each option.
| gathered_buf = pld.alloc_window_buffer([MAT_M, MAT_K], dtype=pl.FP16) | ||
| signal_buf = pld.alloc_window_buffer([world_size, M_TILES], dtype=pl.INT32) | ||
|
|
||
| for rank in pl.range(pld.world_size()): | ||
| gathered = pld.window(gathered_buf, [MAT_M, MAT_K], dtype=pl.FP16) | ||
| signal = pld.window(signal_buf, [world_size, M_TILES], dtype=pl.INT32) | ||
| self.chip_orch(inputs[rank], weight, outputs[rank], gathered, signal, device=rank) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Determine window buffer initialization and lifetime across repeated launches.
set -euo pipefail
rg -n -C10 'def alloc_window_buffer|def window\(' --glob '*.py' || true
# Does any example or test reset a signal window between launches?
rg -n -C6 'signal' examples --glob '*.py' || true
# Benchmark launch loop: confirm the compiled program is re-executed in-process.
rg -n -C15 'def benchmark' --glob '*.py' | rg -n -C15 'rounds|warmup|for ' || trueRepository: hw-native-sys/pypto-lib
Length of output: 161
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== files =="
git ls-files | sed -n '1,120p'
echo "== target file excerpts =="
if [ -f examples/advanced/l3_allgather_gemm.py ]; then
wc -l examples/advanced/l3_allgather_gemm.py
sed -n '1,280p' examples/advanced/l3_allgather_gemm.py
fi
echo "== runtime_benchmark mentions =="
rg -n -C8 'runtime_benchmark|benchmark_rounds|benchmark_warmup|--benchmark-warmup|--benchmark-rounds' .Repository: hw-native-sys/pypto-lib
Length of output: 22642
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate runtime benchmark implementation =="
fd -a 'benchmark.py|*.py$' . | sed -n '1,200p'
rg -n 'def benchmark|class BenchmarkStats|`@benchmark_round`|rounds|warmup' pypto pypto.runtime . --glob '*.py' || true
echo "== target runtime benchmark source snippets =="
rg -n -C80 'def benchmark|class BenchmarkStats|rounds|warmup|compiled' pypto --glob '*.py' || trueRepository: hw-native-sys/pypto-lib
Length of output: 518
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate runtime benchmark implementation =="
git ls-files '*.py' | rg 'pypto|runtime|benchmark' || true
rg -n 'def benchmark|class BenchmarkStats|`@benchmark_round`|rounds|warmup|alloc_window_buffer|window\(' . --glob '*.py' || true
echo "== target runtime paths snippets from all python files =="
for f in $(git ls-files '*.py'); do
if rg -q 'def benchmark|class BenchmarkStats|rounds|warmup|alloc_window_buffer|window\(' "$f"; then
echo "--- $f ---"
wc -l "$f"
rg -n -C12 'def benchmark|class BenchmarkStats|rounds|warmup|alloc_window_buffer|window\(' "$f" || true
fi
doneRepository: hw-native-sys/pypto-lib
Length of output: 50380
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== golden runner relevant region =="
sed -n '920,1085p' golden/runner.py
echo "== runtime benchmark symbol locations =="
rg -n 'def benchmark\(|class Benchmark|rounds=.*warmup|warmup.*rounds|runtime_benchmark|benchmark as' . --glob '*.py' --max-count 200 || trueRepository: hw-native-sys/pypto-lib
Length of output: 9698
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== pypto runtime files =="
git ls-files 'pypto/runtime*' .runtime* benchmarks 2>/dev/null || true
rg -n 'def benchmark\(|def runtime_benchmark|class Benchmark|rounds:|warmup:|persistent|reset_persistent_windows|alloc_window_buffer|window\(' pypto . --glob '*.py' --max-count 300 || trueRepository: hw-native-sys/pypto-lib
Length of output: 50379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== benchmark test region =="
sed -n '1245,1315p' tests/golden/test_runner.py
sed -n '1323,1430p' tests/golden/test_runner.py
echo "== main benchmark call region =="
wc -l examples/advanced/l3_allgather_gemm.py
sed -n '350,420p' examples/advanced/l3_allgather_gemm.pyRepository: hw-native-sys/pypto-lib
Length of output: 9022
Reset signal between benchmark rounds.
The benchmark path reuses the persistent program/window buffers by design, while this kernel leaves signal_buf untouched after each launch and waits with expected=1. Clear the signal window before every repeated benchmark launch, or use round-dependent pld.system.notify values and pld.system.wait thresholds instead.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/advanced/l3_allgather_gemm.py` around lines 234 - 240, Reset the
persistent signal buffer before each repeated benchmark launch in the kernel
containing the all-gather loop and signal window setup. Ensure each round clears
or reinitializes signal values before the `pld.system.wait` checks expecting 1,
without changing the existing gathered buffer or rank orchestration behavior.
1 participant
This PR adds an L3 AllGather-GEMM fusion example that overlaps inter-device communication with GEMM computation at the tile level.
This PR depends on hw-native-sys/PTOAS#1086 and PTOAS 0.56, which has already been released.