Fix: preserve args dump selection in host-built graphs - #1808
Merged
ChaoZheng109 merged 3 commits intoAug 13, 2026
Merged
Conversation
Host-built graph orchestration runs outside AICPU, so its dump-mask and scalar-dtype table updates resolved to host-side no-op stubs. Partial dumps consequently selected no tasks, while hybrid dumps emitted metadata without marked tensor payloads. Embed the selection mask, ambiguity flags, and scalar dtypes in each HBG task payload and cached Graph definition. Consume that metadata in dispatch, completion, and timeout dump paths while retaining the existing AICPU side table for tensormap-and-ringbuffer. Cover direct submission and Graph replay on both architectures, and exercise partial and hybrid manifests on a2a3 onboard plus a2a3sim/a5sim. Fixes hw-native-sys#1800
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughHost-build-graph now embeds per-task argument-dump metadata in task payloads and cached graph nodes. Scheduler dump paths consume this metadata directly. Tests verify selective capture, scalar dtypes, and graph replay for a2a3 and a5. ChangesArgument-dump metadata pipeline
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/common/platform/include/aicpu/args_dump_aicpu.h (1)
368-396: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winFilter unselected tensors in PARTIAL mode.
Lines 368-396 set
capture_payload, but they do not skip unselected tensor arguments.dump_arg_recordonly uses this field to suppress payloads inHYBRIDmode.PARTIALmode therefore records and copies every tensor payload, including unmarked arguments.Proposed fix
ArgsDumpRole role; if (!get_dump_arg_role_from_direction(dir, &role) || !should_dump_arg_at_stage(role, stage)) { tensor_arg_index++; continue; } + const bool capture_payload = has_dump_arg_flag(dump_arg_mask, sig_idx); + if (g_dump_args_level == DumpArgsLevel::PARTIAL && !capture_payload) { + tensor_arg_index++; + continue; + } const auto &t = tensor_info[tensor_arg_index]; ArgsDumpInfo info = {}; @@ - info.capture_payload = has_dump_arg_flag(dump_arg_mask, sig_idx); + info.capture_payload = capture_payload;🤖 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 `@src/common/platform/include/aicpu/args_dump_aicpu.h` around lines 368 - 396, Update the tensor-processing loop around dump_arg_mask, should_dump_arg_at_stage, and capture_payload so PARTIAL mode skips unselected tensor arguments before creating or recording ArgsDumpInfo. Preserve unselected arguments only where the active mode requires metadata without payload, such as HYBRID, and keep tensor_arg_index synchronized for every non-scalar argument.src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp (1)
176-185: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winGuard payload access in both completion paths.
Both completion implementations can dereference a null payload when argument dumping is enabled.
src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp#L176-L185: requireslot_state.payload != nullptrbefore readingdump_metadata.src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp#L176-L185: apply the same payload guard.🤖 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 `@src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp` around lines 176 - 185, Guard the argument-dumping condition in both completion implementations so it requires slot_state.payload to be non-null before accessing dump_metadata. Apply this change in src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp lines 176-185 and src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp lines 176-185, preserving the existing is_dump_args_enabled() behavior otherwise.
🧹 Nitpick comments (1)
tests/ut/cpp/a2a3/test_graph_cache.cpp (1)
335-335: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winCover all static dump metadata during affine replay.
Both graph-cache tests verify only
dump_arg_maskafter affine materialization. A replay-specific loss ofdump_arg_flagsorscalar_dtypescould therefore pass.
tests/ut/cpp/a2a3/test_graph_cache.cpp#L335-L335: assert the first node'sdump_arg_flagsand scalar dtype after affine materialization.tests/ut/cpp/a5/test_graph_cache.cpp#L335-L335: add the same affine replay assertions.🤖 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 `@tests/ut/cpp/a2a3/test_graph_cache.cpp` at line 335, Extend the affine-materialization assertions in tests/ut/cpp/a2a3/test_graph_cache.cpp at lines 335-335 to verify the first node’s dump_arg_flags and scalar dtype alongside dump_arg_mask; apply the same assertions in tests/ut/cpp/a5/test_graph_cache.cpp at lines 335-335, using the expected metadata values established by the test setup.
🤖 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.
Outside diff comments:
In
`@src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp`:
- Around line 176-185: Guard the argument-dumping condition in both completion
implementations so it requires slot_state.payload to be non-null before
accessing dump_metadata. Apply this change in
src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp
lines 176-185 and
src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp lines
176-185, preserving the existing is_dump_args_enabled() behavior otherwise.
In `@src/common/platform/include/aicpu/args_dump_aicpu.h`:
- Around line 368-396: Update the tensor-processing loop around dump_arg_mask,
should_dump_arg_at_stage, and capture_payload so PARTIAL mode skips unselected
tensor arguments before creating or recording ArgsDumpInfo. Preserve unselected
arguments only where the active mode requires metadata without payload, such as
HYBRID, and keep tensor_arg_index synchronized for every non-scalar argument.
---
Nitpick comments:
In `@tests/ut/cpp/a2a3/test_graph_cache.cpp`:
- Line 335: Extend the affine-materialization assertions in
tests/ut/cpp/a2a3/test_graph_cache.cpp at lines 335-335 to verify the first
node’s dump_arg_flags and scalar dtype alongside dump_arg_mask; apply the same
assertions in tests/ut/cpp/a5/test_graph_cache.cpp at lines 335-335, using the
expected metadata values established by the test setup.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 6e5423b6-60e0-4fcd-a25a-c3dff6e67363
📒 Files selected for processing (29)
docs/dfx/args-dump.mdsrc/a2a3/runtime/host_build_graph/host/host_orch_compat_stubs.cppsrc/a2a3/runtime/host_build_graph/runtime/graph_execution.hsrc/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cppsrc/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.hsrc/a2a3/runtime/host_build_graph/runtime/scheduler/graph_execution.cppsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cppsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cppsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cppsrc/a5/runtime/host_build_graph/host/host_orch_compat_stubs.cppsrc/a5/runtime/host_build_graph/runtime/graph_execution.hsrc/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cppsrc/a5/runtime/host_build_graph/runtime/pto_runtime2_types.hsrc/a5/runtime/host_build_graph/runtime/scheduler/graph_execution.cppsrc/a5/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cppsrc/a5/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cppsrc/a5/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cppsrc/common/platform/include/aicpu/args_dump_aicpu.hsrc/common/platform/include/common/args_dump.hsrc/common/platform/shared/aicpu/args_dump_aicpu.cpptests/st/a2a3/host_build_graph/dump_args/kernels/orchestration/dump_args_orch.cpptests/st/a2a3/host_build_graph/dump_args/test_dump_args_example.pytests/st/a5/host_build_graph/dump_args/kernels/orchestration/dump_args_orch.cpptests/st/a5/host_build_graph/dump_args/test_dump_args_example.pytests/ut/cpp/a2a3/test_args_dump.cpptests/ut/cpp/a2a3/test_graph_cache.cpptests/ut/cpp/a2a3/test_hbg_submit_poison.cpptests/ut/cpp/a5/test_graph_cache.cpptests/ut/cpp/a5/test_hbg_submit_poison.cpp
Host-build-graph runtime types only need the dump selection POD, but including the full args_dump header pulled platform timing symbols into orchestration sources and conflicted with paged-attention definitions. Move the task metadata and its ABI assertions into a dependency-light header, then include that header from both a2a3 and a5 runtime types. Paged-attention scene compilation and dump-args level 1/3 simulation tests pass on the supported platforms.
ChaoZheng109
approved these changes
Aug 13, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Host-built graph orchestration runs outside AICPU, so its dump-mask and scalar-dtype table updates resolved to host-side no-op stubs. Partial dumps consequently selected no tasks, while hybrid dumps emitted metadata without marked tensor payloads.
Embed the selection mask, ambiguity flags, and scalar dtypes in each HBG task payload and cached Graph definition. Consume that metadata in dispatch, completion, and timeout dump paths while retaining the existing AICPU side table for tensormap-and-ringbuffer.
Cover direct submission and Graph replay on both architectures, and exercise partial and hybrid manifests on a2a3 onboard plus a2a3sim/a5sim.
Fixes #1800