refactor(EPv2): replace per-module __device__ staging with shared buffer by kawhil-amd · Pull Request #703 · ROCm/mori

kawhil-amd · 2026-09-23T11:26:55Z

Summary

The six __device__ staging arrays in the gfx125x dispatch kernel become one
caller-allocated buffer, passed through EpArgs::stagingBase. The kernel derives
typed pointers at constexpr offsets from the base — same variable names, same
indexing, body otherwise unchanged. One buffer per op instance, shared by every
dispatch schedule variant (they are mutually exclusive). Combine stops carrying
staging it never reads.

There is now exactly one layout. An earlier revision of this PR kept the arrays
as a fallback for callers that never bind the base, selected per launch; that
scaffolding is gone — see Caller contract below, which is the one breaking
part of this change.

Motivation

.bss is reserved per compiled variant, and every staging pool is quadratic
in world size: EpStagingPoolSlots is worldSize * EpMaxRecv, and EpMaxRecv is
itself worldSize * maxTokPerRank. Each JIT variant (a different block/warp
schedule) carried its own copy, and both the dispatch and the combine module
included them, though only dispatch writes them.

At EP32 the scale staging alone tripped a 1 GiB static_assert, so wider EP did
not build at all.

Memory

The whole staging region is one allocation now, sized by ep_staging_bytes() /
EpStagingTotalBytes, with .bss back to a single byte:

$ llvm-readelf -S <gfx1250 dispatch module, EP8 maxTok 128, 224 B scales>
  [10] .bss   NOBITS   000000000000c6f0  0096f0  000001  00  WA  0  0  1

Same for the combine module, and no _cusplit_* symbol in either.

What one op instance allocates, at the 256 B stride a 224 B row (hidden 7168)
pads up to:

Config	staging buffer
EP4 / maxTok 4096	24.3 MiB
EP4 / maxTok 16384	97.0 MiB
EP8 / maxTok 16384	388.0 MiB
EP32 / maxTok 4096	1.52 GiB

against which main pays, per loaded module: ~390 MiB for a dispatch module and
~134 MiB for a combine module at EP8/maxTok 16384 (measured in review on main).
A multi-entry schedule multiplies that; this does not.

EP32/maxTok 8192 compiles here and fails the static_assert on main.

Caller contract — breaking for direct-`EpDispatchPlan` callers

EpDispatchCombineOpHip binds the buffer itself and needs nothing from callers.
A caller that drives EpDispatchPlan directly must now allocate and bind it:

from mori.ops.dispatch_combine_v2.hip_backend import ep_staging_bytes

n = ep_staging_bytes(world_size, max_recv, scale_bytes, max_tok_per_rank)
buf = torch.zeros(n, dtype=torch.uint8, device="cuda")
plan.bind(staging_base=buf.data_ptr())   # keep `buf` alive as long as the plan

aiter's MegaMoE does drive the plan directly, via _build_mori_dispatch. It
needs the bind landed alongside this.

Not binding is a host-side RuntimeError naming the argument, not a memory
fault — EpDispatchSpec::LaunchRaw shadows the base LaunchRaw to reject a null
or under-aligned base, and a device assert backs it up for a caller that names
KernelSpec<...>::LaunchRaw past the shadow.

max_recv is EpCfg.maxRecv, whose 0 means "worldSize * maxTokPerRank"
(EpMaxRecv) and is the Plan default — hence the fourth argument. Passing 0
for both is a ValueError, never a short buffer.

Changes

include/mori/ops/dispatch_combine_v2/ep_cfg.hpp

void* stagingBase in EpArgs, in MORI_EP_ARGS_FIELDS, field count 25 → 26
the constexpr layout: EpStagingTotalBytes, EpStaging*Offset, EpStaging*Bytes

include/mori/ops/dispatch_combine_v2/ep_spec.hpp

EpDispatchSpec::LaunchRaw shadows the base to reject an unbound or
under-aligned staging_base on the host. The sub-offsets are all
EpAlignUp128 multiples of the base, so the base carries the alignment the
arrays used to declare for themselves — 128 B, or every staging TDM run starts
off-row. Gated on plan.geom.sharedBytes, the EpArchIs1250() decision already
taken at Prepare, so portable bodies are untouched.

include/mori/jit/v2/spec.hpp

KernelSpec::Launch calls Derived::LaunchRaw, like Prepare already does for
Geometry and RenderSource. Unqualified it bound to the base, so a spec that
shadows LaunchRaw to validate an argument had the check skipped on that path
while plan_api.hpp's C ABI (SPEC::LaunchRaw) ran it.

src/ops/dispatch_combine_v2/ep_intranode_1250x.hpp

staging pointers derived from args.stagingBase at the top of
EpDispatch1250xBody; static_asserts that CUSPLIT_MAX_* agree with
EpStagingMax*
kEpScaleSlots dropped — it was the extent of the removed array. The
explanation that sat around it is kept next to kEpScaleRows: that is the
per-peer stride and NOT _stgCap, which is smaller and walks into the next
peer's rows.

src/ops/dispatch_combine_v2/ep_spec.cpp

stops emitting MORI_EP_SCALE_SLOTS. Changes the rendered text, so
scaleBytes > 0 configs take one JIT recompile; scaleBytes == 0 renders no
scale defines at all and is unaffected.

python/mori/ops/dispatch_combine_v2/hip_backend.py

ep_staging_bytes(), public, mirroring EpStagingTotalBytes so out-of-tree
callers size the buffer from the same function this op does
allocate dispatch_staging and bind it on dispatch plans, before
_build_kernels (plans bind at construction), and release it in
_close_backend after the plans

The last commit is black + clang-format over the touched files only — the
pre-commit workflow runs each hook over whole files, so this PR inherits the
reflow main already owed on ep_intranode_1250x.hpp. No line changes what it
does; it is separate so the four fixes do not have to be read through it.

Review items from the previous round

Memory saving did not materialise / combine modules carried the arrays —
both followed from keeping the arrays. They are gone.
ep_staging_bytes(max_recv=0) returned a short buffer in silence — fixed,
with the EpMaxRecv fallback and a ValueError where the cap is unknowable.
Pinned by a test.
KernelSpec::Launch bypassed the host check — fixed.
EpStagingTotalBytes's : 0 branch is dead — removed; blkBase/blkCount
are sized independently of the Cfg, so the total is never zero.
EpScaleStride's comment still cited the removed static_assert — rewritten,
along with the "per compiled variant" clause that stopped being true.
kEpScaleSlots / MORI_EP_SCALE_SLOTS unused — removed.
test_jit_core.cpp collided with test(jit): check EpCfg's field count against VisitFields, not a literal #727 — that commit is dropped; test(jit): check EpCfg's field count against VisitFields, not a literal #727 is on
main and makes the same change.

Test plan

Verified on this revision:

gfx1250 cross-compile of the rendered TU (dispatch with and without scales,
and combine): all build, and .bss is 1 B in each with no _cusplit_* symbol
ep_staging_bytes pinned against EpStagingTotalBytes at four shapes, plus
the max_recv=0 resolution and the ValueError
pre-commit clean

Verified in review on the previous revision (4× gfx1250), still expected to hold —
the kernel body did not change between them:

test_op, test_asym_dtype (fp8/fp4 dispatch + 224 B scales at hidden 7168,
topk 6/8/9, ragged), test_jit_binding match main
EP32 maxTok 8192 compiles; main fails the static_assert
bench_ep EP4 bf16/fp8, m 32–4096, eager and graph: within noise of main

Still needed before merge:

re-run the gfx1250 correctness suite on this revision, with no caller
relying on an unbound base
land the aiter-side staging_base bind together with this
gfx942/gfx950 unaffected (portable kernel has no staging) and internode
unaffected (separate staging mechanism)

🤖 Generated with Claude Code

jhchouuu

Tested on 4x gfx1250 (f01-1). The build I tested was e5bc808 with the field count bumped, which matches 6652dcc apart from formatting.

hip intranode test_op.py (EP4: bf16/scatter/topk9/scales/tokoff on+off): same results as main, case for case.
DSv4 e2e (ATOM + aiter MegaMoE, fp4 wire, real routing, gsm8k): faults at startup unless aiter binds staging_base (see inline at L249). With a one-line bind on the aiter side it serves, and gsm8k accuracy matches main.

The layout itself checks out: _ep_staging_bytes matches EpStagingTotalBytes field by field, all staging indexing is size_t (so the >2 GiB scale region at EP32 is fine without the old 1 GiB assert), every staging write is bounds-guarded (dt < _stgCap, dt < kEpScaleRows, blockIdx < 512), and the 128 B sub-offsets on a torch base satisfy the TDM row alignment.

The main blocker is the direct-Plan caller below; the rest are nits.

jhchouuu · 2026-09-28T02:52:00Z

+
+  // Staging pointers from the dynamically-allocated base.  The EpStaging*Offset
+  // functions are constexpr on kCfg, so each is base + compile-time constant.
+  char* _stgBase = static_cast<char*>(args.stagingBase);


Unbound plan args default to 0, so any caller that drives EpDispatchPlan directly gets stagingBase == nullptr and faults here. aiter main's MegaMoE does exactly that on gfx1250 (mega_moe.py _build_mori_dispatch, which bypasses EpDispatchCombineOpHip). With this PR and unmodified aiter, DSv4 e2e dies at startup:

Memory Fault Error [faulting addr: 0x208d000, kernel: mori_ep_dispatch_tdm_fp4x2_ws4_h3584_k6_128x16] Reason: Page not present

The address is exactly null + a staging sub-offset. Unlike tokOffPeers there is no fallback. Please (a) fail fast on the host when staging_base is unbound on gfx125x, not with a GPU fault, and (b) land the aiter-side bind together with this.

jhchouuu · 2026-09-28T02:52:00Z

+
+    Mirrors EpStagingTotalBytes in ep_cfg.hpp.
+    """
+    pool = world_size * max_recv


This hand-mirrors the C++ layout, and direct-Plan callers need it too. To bind aiter I had to import this private helper and redo EpMaxRecv's fallback myself: max_recv == 0 here gives pool = 0 and an undersized buffer, while the kernel uses worldSize * maxTokPerRank. Consider exporting EpStagingTotalBytes(kCfg) through plan.info (or a public helper that applies the same fallback), so the size has one source of truth.

jhchouuu · 2026-09-28T02:52:00Z

+
+constexpr size_t EpStagingTotalBytes(const EpCfg& c) {
+  size_t tail = EpStagingScaleOffset(c) + EpStagingScaleBytes(c);
+  return tail > 0 ? EpAlignUp128(tail) : 0;


nit: tail is always > 0 here (the idx/wt/src/blk regions are never empty), so the : 0 branch is dead. Separately, the EpScaleStride comment in this file (L349-350) still says "Bounded by the static_assert there", but that assert is gone now. It is outside the diff, so I can't comment on it directly.

jhchouuu · 2026-09-28T02:52:00Z

-// Source row vs the stride we lay it down at; they differ by EpScaleStride's pad.
 constexpr int kEpScaleBytes = MORI_EP_SCALE_BYTES;
 constexpr int kEpScaleStride = MORI_EP_SCALE_STRIDE;
 constexpr size_t kEpScaleSlots = (size_t)MORI_EP_SCALE_SLOTS;


nit: kEpScaleSlots (and the MORI_EP_SCALE_SLOTS define ep_spec.cpp still emits) has no users after this change. A one-line note that the scale staging grows as worldSize^2 * maxTok * stride, and that its per-peer stride is kEpScaleRows, not _stgCap, would keep the useful part of the removed comment.

kawhil-amd · 2026-09-28T04:15:29Z

No perf drop compared to the main branch.

jhchouuu

Re-tested d10e32e merged onto main on 4× gfx1250, against main.

Correctness: test_op, test_asym_dtype (fp8/fp4 dispatch + 224 B scales at hidden 7168, topk 6/8/9, ragged, tokoff off) and test_jit_binding match main, both with staging_base bound and with it left unbound (the .bss fallback).
EP32 maxTok 8192 now compiles (main fails the static_assert); an unbound or misaligned base past the budget raises on the host.
bench_ep EP4 bf16/fp8, m 32–4096, eager and graph: within noise of main.

The open point is memory (inline). Also needs a rebase: it conflicts with #727 in test_jit_core.cpp, which makes the same change, so 6adf2ca can be dropped.

The gfx1250x dispatch kernel declared six __device__ global arrays for metadata staging. Each JIT-compiled variant carried its own copy in .bss, and the total grew as worldSize^2 -- hitting a 1 GiB static_assert at EP32. Replace them with one dynamically-allocated buffer passed through EpArgs.stagingBase. The kernel derives typed pointers at constexpr offsets from the base (same variable names, so the body is unchanged). The buffer is allocated once per op instance and shared by all dispatch schedule variants (they are mutually exclusive). Combine never used the staging and no longer carries it. Memory at EP8/maxTok=16384: ~256 MiB (was ~1.5 GiB across 6 modules). EP32 now compiles and runs where it previously failed the static_assert. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

_build_kernels binds staging_base on dispatch plans, so self.dispatch_staging must exist before the call. Also skip allocation on the internode path (it has its own staging). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

A caller that drives EpDispatchPlan directly -- aiter's MegaMoE does, via _build_mori_dispatch, bypassing EpDispatchCombineOpHip -- never binds staging_base, and unbound plan args default to 0. The gfx125x dispatch body then reads staging at null + a sub-offset and dies with a memory fault whose address names neither the argument nor the caller. Unlike tokOffPeers there is no window to fall back to. Two checks, because one does not cover the other: - EpDispatchSpec::LaunchRaw rejects a null base before the launch. Gated on plan.geom.sharedBytes, which is the EpArchIs1250() decision Geometry() already took at Prepare and left on the plan -- the portable body never reads stagingBase, so an ungated check would break every MI300X dispatch. - The kernel asserts too. LaunchRaw only *shadows* the base: the macro's SPEC::LaunchRaw is qualified and finds the override, but KernelSpec::Launch calls LaunchRaw unqualified from the base's scope and gets the base one. No caller does that today; the assert is what keeps it from becoming a fault again if one appears. JIT flags carry no -DNDEBUG, and hip_runtime.h does not pull in assert -- hence <cassert>. Rename _ep_staging_bytes -> ep_staging_bytes so the aiter side can size the buffer from the same function this op does, rather than re-deriving EpStagingTotalBytes (and EpMaxRecv's max_recv==0 fallback) by hand. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

MORI_EP_ARGS_FIELDS is 25 on main; this branch adds stagingBase, so the assert goes to 26. It is deliberately a literal the compiler cannot derive -- the count exists to make adding a field to EpArgs without adding it to the X-macro a build error rather than a binding that writes every later argument into the wrong slot. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

The .bss arrays declared their own alignment -- alignas(kTdmRowBytes) on the idx/wt/src pools, __align__(EpScaleAlign) on the scale rows. The dynamic layout cannot: EpStaging*Offset only places the sub-arrays relative to the base, so the base has to carry it instead. Nothing checked that. A caller carving staging out of a pool at an arbitrary offset -- the reason stagingBase exists -- got no fault, just every staging TDM run starting off-row and degrading to a scalar tail. torch's allocator returns 512 B-aligned pointers, which is why the in-tree path never saw it. Every offset is an EpAlignUp128 multiple, so (base + off) % 128 == base % 128: requiring 128 of the base is both necessary and sufficient to reproduce what the arrays guaranteed. EpScaleAlign rather than kTdmRowBytes because tdm_align.hpp is a device header, and ep_intranode_1250x.hpp already static_asserts the former is a multiple of the latter -- so this can never be the weaker of the two. A null base has remainder 0, so the EpHasStaticStaging fallback above is untouched. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

_close_backend explicitly drops _self_first_state, the plans, _tokoff_ext and _dev_comm, but not dispatch_staging -- which is the largest single device allocation the op owns (12 MiB at EP8/512, ~390 MiB at EP8/16384). A process that builds and closes EP ops for several shapes while keeping the objects referenced held every one of them. After the plans, for the same reason _dev_comm is: they carry staging_base by value, so nothing may still be able to launch when the tensor behind it goes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

ep_staging_bytes() hand-mirrors EpStagingTotalBytes (ep_cfg.hpp) and nothing held it there. EpArgs carries the staging base but no length, so the kernel cannot bounds-check what it was given: a mirror that drifts low is not an error but a silent write past the end of a torch tensor, into whatever the caching allocator handed out next. Four golden numbers, in the style the neighbouring self_first_state_bytes test already uses. The last one has max_recv=7, the only shape among realistic configs where a sub-array size is not already a multiple of 128 -- so it is the only case that exercises the mirror's rounding at all. Values checked against EpStagingTotalBytes for each config. black also reflowed four pre-existing lines in this file: it is not clean under the pinned 25.1.0, and the hook rewrites on commit. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

EpCfg.maxRecv == 0 is not a cap of zero, it is "ask EpMaxRecv", which answers worldSize * maxTokPerRank (ep_cfg.hpp). ep_staging_bytes() took the field at face value, so that spelling gave pool = 0 and a buffer holding only the two world_size-sized block arrays: 16 KiB at EP4/maxTok 16384 where the kernel needs 97 MiB. Nothing catches it afterwards. EpArgs carries the base but no length, so the kernel bounds-checks against the cap it resolved, not against the allocation -- the run writes megabytes past the end of a torch tensor, into whatever the caching allocator handed out next, and reports nothing. In-tree callers pass cfg.effective_max_recv and were never exposed. The function is public precisely because out-of-tree callers are not: they size the buffer from a Cfg whose maxRecv is still the default, which is how this was found while binding aiter's MegaMoE, and LaunchRaw's error message sends them here by name. So take max_tok_per_rank alongside and apply the same ternary. When neither spelling is given the cap is genuinely unknown, and the one thing that must not happen is returning a number for it -- raise instead. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Launch called LaunchRaw unqualified, which binds to KernelSpec's own. A spec that shadows LaunchRaw to check an argument on the host -- EpDispatchSpec does, for staging_base -- had that check skipped for every caller coming through Launch, while plan_api.hpp's C ABI names SPEC::LaunchRaw and did run it. Two entries into the same kernel, disagreeing on whether an argument is validated. Qualify it, as Prepare already does for Geometry, RenderSource, EntryName and SourceDeps: in a CRTP base the derived type is the one being described, so an unqualified name here is a silent "the base's version will do". No behaviour change for a spec that does not shadow LaunchRaw, which is every spec but EpDispatchSpec: the lookup resolves back to this class. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

EpStagingTotalBytes ended in `tail > 0 ? EpAlignUp128(tail) : 0`, but blkBase and blkCount are EpStagingMaxBlocks * worldSize whatever the Cfg says, so the scale region is the only one that can be empty and the tail never is. The branch read as if a config could need no staging at all, which is the one thing a caller sizing a buffer must not conclude. Say what is actually true instead. EpScaleStride's comment still ended "Bounded by the static_assert there", and the 1 GiB assert went with the __device__ arrays; it also called the pool one copy per compiled variant, which stopped being so when it became one allocation per op. Both clauses now describe a build failure that cannot happen, which is worse than saying nothing -- the cost is real, it just lands on VRAM at construction rather than on the compiler. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…nted kEpScaleSlots was the extent of the _cusplit_stgScale __device__ array. The array is a region of the shared staging buffer now, placed by EpStagingScaleOffset and sized by EpStagingScaleBytes, and nothing has read the constant since -- nor MORI_EP_SCALE_SLOTS, which RenderEpSource still emitted for it alone. What the removal lost is the explanation that sat around it, and that part is worth keeping: the per-peer stride is kEpScaleRows and NOT _stgCap, which is the per-peer share of the idx/wt pool and smaller, so indexing this region with it walks into the next peer's rows; and the region is quadratic in world size, which is why it dominates EpStagingTotalBytes. Both now sit next to the constants they are about. The same removal makes the "because the array is at file scope" rationale in ep_spec.cpp wrong -- the macros are still file scope, just for the constants and the kMetaFields #if rather than for an extent. Dropping a #define changes the rendered text, so scaleBytes > 0 configs take one JIT recompile. scaleBytes == 0 renders no scaleDefs at all and is unaffected. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

The pre-commit workflow selects files with --from-ref/--to-ref but runs each hook over the WHOLE file, so a PR that edits one line of a file inherits every reflow already owed on it. ep_intranode_1250x.hpp and hip_backend.py both carry some from main, and this branch now edits both. Whitespace and wrapping only; no line changes what it does. Kept apart from the four fixes so neither has to be read through it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

jhchouuu

Re-checked 814868b; everything from the last round is addressed. On 4x gfx1250 the JIT modules bear out the memory section: a dispatch module's .bss goes from 97.0 MiB before this PR to 513 B, a combine module's from 33.0 MiB to 1 B. A caller that drives EpDispatchPlan directly and binds the buffer runs MLPerf DeepSeek-R1 offline end to end (performance VALID, accuracy 80.65 against the 80.545 threshold).

Two non-blocking suggestions inline.

jhchouuu · 2026-10-12T09:26:20Z

    return (scale_bytes + _SCALE_ALIGN - 1) // _SCALE_ALIGN * _SCALE_ALIGN


+def ep_staging_bytes(


Non-blocking: this mirror is held to EpStagingTotalBytes only by the constants in test_staging_bytes_matches_the_device_layout, so a C++ layout change that forgets this function still passes that test. Since EpArgs carries no length, the result would be a silent write past the end of the caller's tensor. Publishing EpStagingTotalBytes(kCfg) through plan.info would leave one source of truth, and a caller that drives EpDispatchPlan directly could size the buffer from the plan it already holds.

jhchouuu · 2026-10-12T09:26:20Z

 #define CUSPLIT_POOL_SLOTS (MORI_EP_WORLD_SIZE * MORI_EP_MAX_RECV)
 #define CUSPLIT_MAX_BLOCKS 512
 #define CUSPLIT_MAX_TOPK 16
+static_assert(CUSPLIT_MAX_TOPK == EpStagingMaxTopk,


Non-blocking, and not new in this PR: nothing bounds numExpertPerToken by EpStagingMaxTopk or blockNum by EpStagingMaxBlocks. EpCfgIsValid admits topk up to waveSize - 1 and, with aggReserve, blockNum < 4096. Topk 17..31 runs a destination's idx/wt rows into the next one (row = d * _stgCap * CUSPLIT_MAX_TOPK + dt * topk), and blockNum > 512 runs past blkBase/blkCount (indexed blockIdx.x * npes + p). Now that the layout lives in ep_cfg.hpp, a static_assert(kCfg.numExpertPerToken <= EpStagingMaxTopk && kCfg.blockNum <= EpStagingMaxBlocks) in the dispatch body would turn that into a build error.

…opy it ep_staging_bytes() reproduced the staging layout in Python, and nothing held the two together: the test pinned the copy to constants, which go stale alongside the function they pin. A C++ layout change that never reached hip_backend would pass it -- and EpArgs carries the base but no length, so the kernel cannot bounds-check what it was given. A copy that drifts low is a silent write past the end of the caller's tensor, not a fault. The layout cannot move to Python: the kernel derives its six staging pointers from EpStaging*Offset(kCfg) as compile-time constants, so C++ is where it is consumed, and EpStagingTotalBytes is just the last offset plus the last size over those same functions. What was optional was how Python learns the number. So ask. libmori_ops_v2.so exports mori_ep_staging_bytes, ep_plans binds it beside the plans it already loads, and ep_staging_bytes keeps its signature and its ValueError but forwards the rest. maxRecv == 0 now resolves through EpMaxRecv on the C++ side, the same call the kernel's indexing goes through. The symbol sits outside the generic plan ABI deliberately. plan_api stays op-agnostic -- it only grows a return value on load_library, so an op-library carrying symbols of its own can bind them without repeating the search -- and the EP-specific binding lives in ep_plans with the rest of EP's naming. A plan reports no size: that would mean compiling a kernel to answer "how many bytes", and a caller that allocates before it builds a plan could not use it. test_jit_binding's four golden values and both EpMaxRecv cases are unchanged and still pass -- they pin the C++ layout now rather than a copy of it. Checked by calling the export through ctypes, and against EpStagingTotalBytes directly. Raised in review on ROCm#703 (jhchouuu), twice. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…an replace (#764) #703 made the gfx125x dispatch read its staging from a caller-bound staging_base. A caller that drives EpDispatchPlan directly and binds none, such as aiter main's MegaMoEGfx1250 on ATOM's EP+DPA path, failed at its first dispatch. - An EpDispatchPlan that stages (sharedBytes > 0) starts with a default staging_base. A caller's bind() replaces it before it is ever read. - Left unbound, the default resolves at the first launch to an ep_staging_bytes() buffer shared by the unbound plans of one layout on the device. It is then cached as an int, like a bound address. A FutureWarning names the caller's line. - Allocation is refused inside a graph capture; a plan first launched there reuses the buffer its layout already has. - A launch group compares the default without resolving it. - plan_api: a launch argument wins over a non-int bound default on every launch, and LaunchGroup compares pointer-like defaults by address.

kawhil-amd requested review from jhchouuu and zhangfei829 and removed request for jhchouuu September 23, 2026 11:28

kawhil-amd force-pushed the refactor/ep-staging-shared-buffer branch 3 times, most recently from 1863dab to 6652dcc Compare September 24, 2026 08:27

jhchouuu reviewed Sep 28, 2026

View reviewed changes

kawhil-amd force-pushed the refactor/ep-staging-shared-buffer branch 3 times, most recently from 1b68ff1 to d10e32e Compare October 8, 2026 10:13

jhchouuu reviewed Oct 10, 2026

View reviewed changes

kawhil-amd and others added 12 commits October 12, 2026 06:00

kawhil-amd force-pushed the refactor/ep-staging-shared-buffer branch from d10e32e to 814868b Compare October 12, 2026 06:15

jhchouuu reviewed Oct 12, 2026

View reviewed changes

jhchouuu mentioned this pull request Oct 12, 2026

[FlyDSL] [CI] [MoE] MegaMoE gfx1250: stage1/stage2 switches, mori combine path, current mori EPv2 ROCm/aiter#6430

Open

jhchouuu approved these changes Oct 12, 2026

View reviewed changes

kawhil-amd merged commit f839603 into ROCm:main Oct 12, 2026
15 checks passed

kawhil-amd mentioned this pull request Oct 12, 2026

chore(EPv2): bound the staging extents + call EpStagingTotalBytes instead of copying it #752

Open

6 of 10 tasks

srinivamd mentioned this pull request Oct 12, 2026

[EPv2] staging_base binding required by #703 breaks callers that go through aiter MegaMoE (ATOM EP4+DPA) #757

Open

jhchouuu mentioned this pull request Oct 13, 2026

fix(EPv2): give a dispatch plan a default staging buffer its caller can replace #764

Merged

Navigation Menu

Repository navigation

refactor(EPv2): replace per-module device staging with shared buffer - #703

refactor(EPv2): replace per-module device staging with shared buffer#703
kawhil-amd merged 12 commits into
ROCm:mainfrom
kawhil-amd:refactor/ep-staging-shared-buffer

kawhil-amd commented Sep 23, 2026 •

edited

Loading

Uh oh!

jhchouuu left a comment

Uh oh!

jhchouuu Sep 28, 2026

Uh oh!

jhchouuu Sep 28, 2026

Uh oh!

jhchouuu Sep 28, 2026

Uh oh!

jhchouuu Sep 28, 2026

Uh oh!

kawhil-amd commented Sep 28, 2026

Uh oh!

jhchouuu left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jhchouuu left a comment

Uh oh!

jhchouuu Oct 12, 2026

Uh oh!

jhchouuu Oct 12, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		return (scale_bytes + _SCALE_ALIGN - 1) // _SCALE_ALIGN * _SCALE_ALIGN


		def ep_staging_bytes(

Conversation

kawhil-amd commented Sep 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Motivation

Memory

Caller contract — breaking for direct-EpDispatchPlan callers

Changes

Review items from the previous round

Test plan

Uh oh!

jhchouuu left a comment

Choose a reason for hiding this comment

Uh oh!

jhchouuu Sep 28, 2026

Choose a reason for hiding this comment

Uh oh!

jhchouuu Sep 28, 2026

Choose a reason for hiding this comment

Uh oh!

jhchouuu Sep 28, 2026

Choose a reason for hiding this comment

Uh oh!

jhchouuu Sep 28, 2026

Choose a reason for hiding this comment

Uh oh!

kawhil-amd commented Sep 28, 2026

Uh oh!

jhchouuu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jhchouuu left a comment

Choose a reason for hiding this comment

Uh oh!

jhchouuu Oct 12, 2026

Choose a reason for hiding this comment

Uh oh!

jhchouuu Oct 12, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

kawhil-amd commented Sep 23, 2026 •

edited

Loading

Caller contract — breaking for direct-`EpDispatchPlan` callers