Skip to content

Support: restore selected args payloads in L0 replay #1567

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
indigo1973 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Support: restore selected args payloads in L0 replay #1567

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
83 changes: 75 additions & 8 deletions .claude/skills/l0-swimlane/SKILL.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
---
name: l0-swimlane
description: Produce an L0 (intra-core AICore pipeline) swimlane for one task — a single kernel or a mix — via the dump-driven `simpler_setup.tools.l0_swimlane` tool. Use when the user asks to "run/produce an l0 swimlane", "trace a task's intra-core pipeline", profile why one AICore task is slow inside the core(s), or needs help choosing the tool's manual flags (`--func-id`, `--set-arg`, `--spmd-block-num`, `--case`). The tool captures real per-task args from an args dump and auto-generates the `msprof op simulator` replay — no hand-authored workspace. For a hand-authored single-`kernel_entry` replay use [insight-trace](../insight-trace/SKILL.md); for cross-task / scheduler / dependency timing use the L2 swimlane.
description: Produce an L0 (intra-core AICore pipeline) swimlane for one task — a single kernel or a mix — via the dump-driven `simpler_setup.tools.l0_swimlane` tool. Use when the user asks to "run/produce an l0 swimlane", "trace a task's intra-core pipeline", profile why one AICore task is slow inside the core(s), or needs help choosing the tool's manual flags (`--func-id`, `--restore-arg`, `--set-arg`, `--spmd-block-num`, `--case`). The tool captures real per-task args from an args dump and auto-generates the `msprof op simulator` replay — no hand-authored workspace. For a hand-authored single-`kernel_entry` replay use [insight-trace](../insight-trace/SKILL.md); for cross-task / scheduler / dependency timing use the L2 swimlane.
---

# L0 Swimlane — Intra-core Pipeline Trace for a Task
Expand All @@ -10,8 +10,9 @@ reconstructs them, generates a combined `msprof op simulator` replay of the
**whole task** (a mix runs AIC + AIV0 + AIV1 in one op), and exports an
Insight `trace.json` whose lanes are the cluster's pipes. Full reference:
[docs/dfx/l0-swimlane-profiling.md](../../../docs/dfx/l0-swimlane-profiling.md).
This skill is the **operating procedure** — above all the one genuinely
manual decision: the slot/value for `--set-arg`.
This skill is the **operating procedure** — above all the manual decisions:
whether a structured control tensor needs `--restore-arg`, and the
slot/value for any `--set-arg`.

## When to use

Expand Down Expand Up @@ -42,6 +43,15 @@ Onboard `a2a3` instead of `a2a3sim`: run
runs on the locked device). The five internal steps and all flags are in
doc §3.2 / §3.3.

**Payload-restoration scope:** `--restore-arg` is supported only with
`--platform a2a3sim` or `a2a3`. On `a5sim` / `a5`, tensor payload bytes from
args dump are not trustworthy because of
[#1560](https://github.com/hw-native-sys/simpler/issues/1560), so do not use
them as tensor truth or pass `--restore-arg`. Metadata-only and zero-filled
A5 L0 replays remain usable. A5 still writes the selected level-3 `args.bin`,
but an L0 trace restored from it is untrustworthy until #1560 is fixed; this
limitation does not make the whole L0 tool a2a3-only.

## Choosing the manual flags (the hard part)

### `--func-id` — the task's member set
Expand All @@ -53,6 +63,25 @@ lanes SPMD mix the dump records a duplicate (`[0,1,1]`), so pass
`--func-id 0,1` (`set([0,1,1]) == {0,1}`). Wrong set → the tool lists the
func_id shapes present in the dump; pick one of those.

### `--restore-arg SLOT` — real structured control data

This is currently an **a2a3/a2a3sim-only** capability. Do not use A5 payload
for restoration while
[#1560](https://github.com/hw-native-sys/simpler/issues/1560) is open.

Use this repeatable flag when a tensor contains non-uniform runtime-built
state that changes control flow or addresses, such as a tiling struct.
The fresh capture stays at `--dump-args 3`, but adds the internal SceneTest
selector `--dump-args-payload-func-id <FUNC_SET>` plus one
`--dump-args-payload-slot <SLOT>` per restore. It restores those selected
`before_dispatch` bytes and leaves all unselected tensor payloads zero. A
reused `--dump-json` must have its sibling `args.bin` and the requested
record payload.

Do not use it for weights, activations, or KV pools. The selective level-3
capture avoids dumping unrelated task payloads. A restored slot cannot also
use `--set-arg`.

### `--set-arg SLOT=VALUE` — only when a loop count must shrink

First classify where the kernel's loop trip count comes from:
Expand All @@ -62,6 +91,7 @@ First classify where the kernel's loop trip count comes from:
| **Tensor shape** (e.g. `shapes[0] / TILE_ELEMS`) | **No** | shape is the real dump value; changing it distorts. (mixed_example / single-kernel rows need no `--set-arg`.) |
| **A scalar arg** (e.g. `n_blocks`) | **Yes** — set the count directly | camodel would run the full loop; shrink to ≥ 3–4 (doc §7.2: floor 3, prefer 4). |
| **A control-tensor's content** (e.g. `context_lens`) | **Yes** — fill the buffer | the kernel *derives* the count from the data; fill so the derived count ≈ 4 (need `block_size` to back out the value). Integer dtypes only. |
| **A structured control tensor** (e.g. `FAInferTilingData`) | **No** — use `--restore-arg` on a2a3/a2a3sim | One uniform integer cannot reproduce a struct; restore its captured bytes on the supported platform family. |
| **The SPMD `block_num`** | **No** — use `--spmd-block-num` | block_num lives in the synthesised slot-48 context, which `--set-arg` cannot reach. |

Then find the slot — it is **per-kernel, never fixed**. Discover it:
Expand All @@ -86,10 +116,18 @@ Verified examples (slots read from source):
### `--spmd-block-num N` — SPMD grid width

`block_num` is written into the synthesised slot-48 `LocalContext`. Default
is 1; pass it to replay an SPMD cohort, or for a kernel that branches or
grid-strides on `block_num` (e.g. set the real hw width `24`). `block_idx`
is always synthesised to `0` (a representative block) and is **not** a flag —
it has no instruction-stream branches (doc §8).
is 1, which models a single-block replay and does not take multi-block
branches such as `block_idx + 1 < block_num`. Keep that default only when the
kernel ignores `block_num`. For a kernel that branches or grid-strides on it,
read the selected task's real logical width from its orchestration and pass it
explicitly. That value may come from `rt_available_cluster_count()`, or it may
be task-specific. For example, the default lowest task-id in
`spmd_multiblock_aiv` uses `--spmd-block-num 4`, while the default lowest task
in `spmd_multiblock_mix` uses `--spmd-block-num 2`. When the real value is
unavailable, a value of at least 2 is only a bounded approximation for
exercising the next-block branch; record that limitation in the analysis.
`block_idx` is always synthesised to `0` (a representative block) and is
**not** a flag — it has no instruction-stream branches (doc §8).

### `--case NAME` — pin a small case on a multi-case test

Expand All @@ -111,6 +149,35 @@ per-block pipeline stays identical to production (you lose only iteration
*count*, which does not change the pipeline shape). A case with different
tile shapes traces only itself, not production.

### Oversized CAModel checklist

On the supported a2a3/a2a3sim path, selective level-3 payload capture only
shrinks `args.bin`; the replay still allocates and zero-initializes every
tensor at its dumped shape. For a large weight or KV tensor:

1. Run with `--dump-json <manifest> --no-collect` first and inspect the printed
slot shapes.
2. Use `--case <bounded-case>` to reduce the actual shape-driving extents
(layers, pages, cache rows, or sequence blocks). Lowering runtime `seq_len`
alone does not help when shapes use fixed `MAX_SEQ` or layer constants.
3. Use `--set-arg` only when the remaining cost is a loop count; it does not
resize tensors. On a2a3/a2a3sim, use `--restore-arg` only to preserve
structured control state; it does not make the replay smaller. Do not use
A5 payload restoration while #1560 is open.
4. Keep the default `--msprof-timeout 120` for ordinary cases. For a large,
repetitive workload where a partial pipeline is sufficient, set a shorter
timeout such as `--msprof-timeout 4`. This is the behavior recommended by
the [Ascend CANN official QA/parameter guidance for msProf `--timeout`](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/82RC1alpha001/devaids/optool/atlasopdev_16_0082.html):
msProf terminates simulation at the limit and parses the data produced so
far. A partial trace can show the observed pipeline segment, but cannot
establish whole-kernel duration or tail completeness. If the target kernel
produced no profiler data before timeout, there may be no `trace.json`.

If no bounded case exists, make a focused workload-local case with the same
target task, func-id set, argument order, tiling path, and tile geometry but
smaller shape-driving extents. Keep a diagnostic-only case out of the PR unless
it is intended as permanent regression coverage. Full details are in doc §3.6.

## Self-check after every run

A known msprof/camodel export bug can truncate the last loop iteration(s).
Expand All @@ -123,4 +190,4 @@ raw Insight `trace.json`, for sub-laned per-instruction overlap (doc §3.4).

Representative command per task shape (single AIC / single AIV / 1+1 mix /
2-AIV mix / 3-way mix / SPMD single-source / SPMD coop mix / same-AIV-both-
lanes / paged-attn scalar & control-tensor loops / qwen3) is in doc §3.7.
lanes / paged-attn scalar & control-tensor loops) is in doc §3.8.
17 changes: 16 additions & 1 deletion conftest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,7 +156,22 @@ def pytest_addoption(parser):
default=0,
help="Dump per-task args at runtime. Level: 0=off, 1=partial (only "
"tasks marked via Arg::dump(...), default when given without a value), 2=full (all tasks), "
"3=full_json_only (all tasks, JSON metadata only, no .bin payload).",
"3=full_json_only (all tasks' JSON metadata; no payload unless a selective payload filter is given).",
)
parser.addoption(
"--dump-args-payload-func-id",
action="store",
default=None,
metavar="ID[,ID...]",
help="With --dump-args 3, capture tensor payload only for tasks whose active func-id set matches.",
)
parser.addoption(
"--dump-args-payload-slot",
action="append",
type=int,
default=[],
metavar="SLOT",
help="With --dump-args 3, capture this tensor args[] slot for the selected task. Repeatable.",
)
parser.addoption(
"--enable-dep-gen",
Expand Down
Loading
Loading