Repository navigation
Optimize: pack HBG Graph replay data into one arena block - #1923
sunkaixuan2018 wants to merge 3 commits into
Conversation
|
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:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (18)
Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review. 📝 WalkthroughWalkthroughGraph submissions now contain only graph metadata. Boundary arguments reside in outer task payloads. Definitions and submissions are packed into one arena-bank GraphBlock and uploaded in one operation. Host APIs and runner cleanup now manage reusable per-bank GraphBlocks. ChangesGraphBlock submission migration
Estimated code review effort: 4 (Complex) | ~45 minutes Merge Risk: 🔵 Low · up to The PR changes graph replay preparation to use one packed upload. It is mergeable with owner awareness that the packed submissions retain borrowed data pointers across the preparation loop; if that storage is temporary, uploads could read invalid graph data. Sequence Diagram(s)sequenceDiagram
participant graph_submit_outer
participant upload_graph_block
participant HostApi
participant DeviceRunnerBase
participant GraphExecution
graph_submit_outer->>graph_submit_outer: write boundary tensors and scalars to outer payload
upload_graph_block->>HostApi: acquire_graph_block(arena_bank, bytes, alignment)
HostApi->>DeviceRunnerBase: acquire graph-block storage
upload_graph_block->>DeviceRunnerBase: upload packed definitions and submissions
GraphExecution->>graph_submit_outer: read boundary arguments from outer payload
GraphExecution->>GraphExecution: localize and materialize graph execution
Poem 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 |
ChaoWao
left a comment
ChaoWao
left a comment
There was a problem hiding this comment.
The GraphBlock half of this looks right to me, and it goes further than the batching
I had staged locally — one grow-only block on the arena bank with base + offset
addressing and reuse-when-large-enough is the shape the "one rtMalloc, one
rtMemcpy" constraint asks for, where mine only collapsed the per-submission
allocations into a single fresh one per run. I'll drop mine.
One problem in the other half, though: graph_submit_outer writes the boundary
args into payload.tensors / payload.scalars with no bound check, and those arrays
are four times smaller than what GraphTaskArgs can legally hold.
// pto_orchestrator.cpp, this PR
payload.tensor_count = args.tensor_count();
for (int32_t i = 0; i < args.tensor_count(); ++i)
payload.tensors[i] = args.tensor(i).ref();
std::memcpy(payload.scalars, args.scalar_data(), args.scalar_count() * sizeof(uint64_t));| capacity | where | |
|---|---|---|
GraphTaskArgs (the source) |
128 tensors / 64 scalars | pto_types.h:689-690 |
PTO2TaskPayload.tensors / .scalars (the destination) |
32 / 16 | CORE_MAX_TENSOR_ARGS / CORE_MAX_SCALAR_ARGS |
The only cap on the path in between is a debug_assert(args.tensor_count() <= GRAPH_MAX_TENSOR_ARGS) in rt_submit_graph_impl, i.e. 128.
examples/a2a3/host_build_graph/deepseek_v4_flash_decode exceeds both. Measured on
a2a3 by logging the counts at the write site and running the case's host path only
(SIMPLER_SKIP_DEVICE_RUN=1), on this PR's head:
[PROBE] graph boundary: tensors=118 (payload holds 32) scalars=31 (payload holds 16)
[PROBE] would write tensors[0..117] into a 32-slot array: bytes [640, 15744)
against sizeof(payload)=4864, i.e. 10880 bytes past the payload
(2.24 following payloads at stride 4864)
40 occurrences per run — 2 ranks x 20 submissions. tensors starts at offset 640 and
scalars at 4736, so tensors[32..117] runs straight over the scalar region and then
10,880 bytes past the end of a 4,864-byte payload, into the following entries of the
payload array. The counts come from the case's own orchestration
(decode_fwd_graph.cpp:8788-8919: 113 add_input + 5 add_inout, and 31 scalars
across 8 add_scalar calls), so they are static, not workload-dependent. My probe
returned early instead of performing the write, so nothing above is an observed
corruption — it is the extent of the write the code would make.
main also states the invariant this breaks, in the comment on the type being read
here:
Boundary arguments of a Graph. Sized independently of CoreTaskArgs because a Graph
boundary never becomes a task payload —graph_reset_outer_payloadzeroes the outer
GRAPH task'stensor_count, and the boundary is read host-side only […] Widening it
costs orchestration stack, notPTO2TaskPayload[…]
—pto_types.h:692-696
That paragraph would need to change with this PR either way; right now the tree would
carry a comment that contradicts its own code, and GRAPH_MAX_TENSOR_ARGS = 128
becomes a number no caller can actually use.
Four things in fairness, because I do not think this is a regression of a green test:
- dsv4 does not pass on
maineither. Its device path fails with
sched_error_code=5 INVALID_ARGS("an orchestration API rejected its arguments").
So the change does not turn a passing case red; it turns an already-failing case
from a clean rejection into an out-of-bounds host write whose bytes then ship H2D. - It is not silent-wrong-results.
graph_execution_localizedoes gain
tensor_count > MAX_TENSOR_ARGS → return nullptr, so the replay is refused. But
that check runs on device, after the host has already written past the array, so the
reported failure point and the actual fault are one step apart. - The validation set cannot see it. qwen3-14b's boundary is 26 tensors / 0 scalars
and fits in 32/16, andrecord_then_replay_1d/2dare smaller still. dsv4 is
manual: True, so the defaultpytest examples tests/stsweep excludes it — CI will
stay green on this indefinitely. - The benefit and the bug share a root. The 232,312 -> 131,832 B (-43%) saving comes
from deleting the duplicate variable-size arg image, and that is exactly the
mechanism a >32-arg boundary needs. So it is a real design tension, not an oversight
to patch with a clamp: bounding the boundary atMAX_TENSOR_ARGSwould cap every
Graph at 32 boundary args and exclude dsv4 outright.
What I have locally for this is a pooled GraphBoundaryImage — a POD
{total_bytes, tensor_count, scalar_count, off_tensors, off_scalars} header in a
host-written slot of a boundary pool, addressed base + index x stride with the index
being the outer task's submission order, bounded by GRAPH_MAX_* rather than
MAX_TENSOR_ARGS. It keeps the args out of GraphSubmission (so the fixed-size
submission object this PR introduces stays fixed-size) without putting them in a
32-slot array, and the pool's base and stride can live in the same GraphBlock this PR
adds. Happy to rebase it onto this and hand it over, or to just supply the bound check
and a failing test if you would rather keep the scope here as-is.
Reproduction, if useful:
# log args.tensor_count()/scalar_count() at the write site in graph_submit_outer, then:
SIMPLER_SKIP_DEVICE_RUN=1 SIMPLER_LOG_LEVEL=TIMING task-submit --device auto --device-num 2 --run \
"python examples/a2a3/host_build_graph/deepseek_v4_flash_decode/test_deepseek_v4_flash_decode.py \
-p a2a3 --manual only -d \$TASK_DEVICE \
--case TestDeepseekV4FlashDecodeHostBuildGraph:: --runtime host_build_graph --level 3"Note that SIMPLER_SKIP_DEVICE_RUN is presence-based, and that a skip_golden case
reports PASSED under it having executed nothing — so read the probe output, not the
exit status.
- Store boundary tensors and scalars in the outer GRAPH task payload\n- Keep per-replay device state in a fixed-size submission object\n- Remove the duplicate variable-size argument wire image
Compute Definition and submission offsets before acquiring device storage. Store a grow-only GraphBlock on the selected arena bank, reuse its capacity, and overwrite the complete live range with one H2D copy per run.
2a1a0f2 to
1b4278b
Compare
TaskPayloadSpace follow-upAddressed the DSV4 boundary-argument overflow identified in review without changing
Validation:
DSV4 device execution is a known-broken path and is not an acceptance criterion for this PR. Updated head: |
|
Closing as superseded by #1952 ( To be clear about what this PR did achieve, because I measured it rather than inferring it. On a2a3, qwen3-14b decode, its head
What it left in place, also measured on that head:
Why #1952 supersedes it: instead of making the payload variable-size to fit a wide Graph boundary, it moves the variable-extent regions (fanin, tensors, scalars) out of the payload into three pools at the image tail, each shipping only its live prefix, with the payload naming them by delta. That subsumes the arg de-duplication, keeps One item from this line of work is still not in Thanks — the arg de-duplication was the right observation, and it is in |
3 participants
Summary
Reduce host-side allocation and H2D-copy overhead for host-built Graph replay data.
GRAPHtask payload, leavingGraphSubmissionas a fixed-size Definition reference plus device replay state.[Definitions][Submissions]layout, including every size, alignment, and offset, before acquiring device storage.GraphBlock base + offsetand overwrite the complete live range with exactly one H2D copy on every run.This deliberately does not add content-based copy skipping, Task Window packing, queue sizing, heap sizing, or environment-variable-driven allocation.
Motivation
Host
rtMallocis expensive, and one largertMemcpyis materially cheaper than many small copies. The previous path retained each Definition separately and allocated/uploaded every submission independently. Qwen3 performs 40 Graph replays, producing 41 allocation/H2D operations per run.The new path computes the whole GraphBlock first, performs at most one allocation/grow operation for the selected arena bank, and then performs one H2D upload.
Performance
Measured with
GraphExecutionBatch16Seq3500from the Qwen3-14B decode HBG example on A2A3. Baseline and PR were rebuilt from adjacent runs against currentmain(93adc386).graph_uploadhost spanThe direct
graph_uploadspan is the relevant result; end-to-end device execution is expected to remain unchanged because this PR only changes host preparation and upload.Validation
main(93adc386).record_then_replay_1drecord_then_replay_2d(multiple Definitions and submissions)GraphExecutionBatch16Seq3500correctness and profiling passed.A5 onboard hardware was not available on the A2A3 validation host; A5 compilation, C++ tests, and A5 simulator coverage passed.