[Feature] In-kernel put-fence / TPUT-quiesce primitive to enable multi-core parallelized cross-rank push (combine_push / dispatch)

[lwDavid]

Summary

The MoE combine_push (and the symmetric dispatch payload push) in the DeepSeek-V4 kernels is a serial, put-issue-bound task that cannot be parallelized across cores today because pypto has no in-kernel primitive to guarantee a rank's outbound TPUTs have landed at the peer before a subsequent cross-rank notify.

We need either:

• (A) an in-kernel put-fence / TPUT-quiesce op (per-rank drain of outbound TPUTs), usable mid-kernel between a multi-core push scope and a single-core barrier scope; and/or
• (B) formal support + validation for a multi-core pl.spmd cross-rank pld.tensor.put followed by a deps-ordered single-core notify/wait barrier, with the semantics that the deps edge implies remote-landing (not just task-completion).

Motivation / Use Case

combine_push (pypto-lib/models/deepseek/v4/combine.py:92-153) is one pl.at(CORE_GROUP) serial task: one scalar core builds ~48 scattered pld.tensor.put descriptors (~0.5 us each), then a cross-rank combine_done notify/wait barrier. Measured ~24 us on a2a3 EP=2, and during its window all 71 other cores are idle.

It is put-ISSUE-bound, not transfer-bound:

• 384 KB / 24 us = 16 GB/s ≈ 17% of the ~96 GB/s link (6× headroom) — a transfer-bound kernel would sit near peak.
• 24 us / 48 puts ≈ 0.5 us/put = one scalar core serially building scatter descriptors. Scattered dst_offsets=[r_route, 0] forbid coalescing, so each row is its own descriptor.

The hardware already supports parallel issue: each core's TPUT selects channelGroupIdx = get_block_idx() → its own SDMA SQ/CQ among kSdmaMaxChannel = 48 channels. Distributing the put-issue across the idle cores would cut combine_push to a ~5-6 us floor (transfer + barrier), saving ~16-18 us per call. Because it is on the fully-serial MoE tail, that converts 1:1 to wall-clock and multiplies across the model's MoE layers per decode step.

Proposed API / Behavior

Option A — explicit in-kernel put-fence (preferred; minimal new surface):

with pl.spmd(EP_WORLD_SIZE * N_LOCAL_EXPERTS, name_hint="combine_push") as push_tid:
    bucket = pl.tile.get_block_idx()
    ...                                  # each core issues its own bucket's puts on its own SDMA channel
    pld.tensor.put(dst=routed_y_buf, peer=dst, src=recv_y_flat, ...)

with pl.at(level=pl.Level.CORE_GROUP, deps=[push_tid], name_hint="combine_done") as barrier_tid:
    pld.system.put_fence()               # NEW: block until ALL of this rank's outbound TPUTs have LANDED at peers
    pld.system.notify(target=combine_done, ...)   # now guaranteed to follow the data
    pld.system.wait(signal=combine_done, ...)

Option B — make deps imply TPUT-landing: define that a scope with deps=[push_tid], where push_tid issued cross-rank puts, is guaranteed to start only after push_tid's outbound TPUTs have landed (not merely been issued to SDMA), and add an ST test mirroring this combine_push pattern.

Either way the deliverable is: a kernel author can split put-issue across cores and still get correct cross-rank ordering before the done-notify.

Alternatives Considered

• Flat static-bound form (serial route-table build + pl.spmd over N_ROUTES with a per-lane mask): retires the dynamic-loop-bound 507018 class, but keeps the same push/barrier split → same cross-core ordering hazard → expected to deadlock identically. Not a fix.
• Per-expert streaming (push inside expert_routed's per-expert pl.parallel to hide under the expert cloud): needs the exact same cross-core put→notify ordering guarantee plus an all-experts join, with larger blast radius — strictly worse than fixing the primitive.
• Keep push+notify atomic (status quo): correct, but leaves ~16-18 us/layer of put-issue serialized on one core with 71 cores idle.

Additional Context

Empirical result. We implemented the 3-scope split (with pl.spmd(EP*N_LOCAL) as push_tid → pl.at(CORE_GROUP, deps=[push_tid]) barrier → pl.spmd(deps=[barrier_tid]) reduce, capture-form with pl.tile.get_block_idx()). It compiles, but on a2a3 EP=2 it DEADLOCKS (300 s timeout, killed) on the first run. This multi-core cross-rank put + deps-separated barrier pattern is not exercised anywhere in pypto tests.

Why it breaks. The current code deliberately keeps the structure atomic — combine.py:88: "keep TPUT + notify/wait in one atomic task". In one InCore task, the puts and the notify share the same SDMA channel + program order, so the notify is guaranteed to follow the puts. After the split, the barrier's notify runs on a different SDMA channel than the puts (channelGroupIdx = get_block_idx()), with no hardware ordering between this rank's puts-landing and its notify-to-peer. deps=[push_tid] only guarantees task-completion of the push scope, not remote-landing of its TPUTs. There is no in-kernel "drain my outbound TPUTs" primitive (comm_barrier is host-side all-rank, unusable mid-kernel).

Hardware path references: per-core channel selection channelGroupIdx = get_block_idx() (pto-isa .../comm/async/sdma/sdma_async_intrin.hpp:514-532), kSdmaMaxChannel = 48 (sdma_types.hpp:30); TPUT emit is pto.barrier<PIPE_ALL> + pto.comm.tput with no post-put completion/quiesce barrier (src/backend/common/pto_ops_common.cpp:2913-2919).

Affected kernels: pypto-lib/models/deepseek/v4/combine.py:92-153 (combine_push), dispatch.py:209-235 (symmetric payload push).

Versions: pypto 55fdd38 (main), pypto-lib 1a23bd3. Target: a2a3 EP=2. Blocks a measured ~16-18 us/layer MoE-tail win.

[lwDavid]

Update — sharper diagnosis after on-device experiments (the put-fence/ordering framing was NOT the blocker)

Tested the parallel combine_push directly on a2a3 EP=2 across several structures. Key correction to the original framing: co-locating put and notify in the same pl.spmd lane is accepted by codegen (each lane keeps its put+notify on the same core/SDMA channel = the original ordering guarantee), so the cross-core put→notify ordering is not the real blocker. Two concrete blockers surfaced instead:

1. Runtime blocker (the real one) — multi-core spmd cross-rank put/notify drain-fails

Structure tested (counting barrier so each lane self-notifies, no cross-core fence needed):

with pl.spmd(EP_WORLD_SIZE * N_LOCAL_EXPERTS, name_hint="combine_push") as push_tid:
    bucket = pl.tile.get_block_idx(); dst = bucket // N_LOCAL; e = bucket - dst*N_LOCAL
    ... for row in pl.range(n): pld.tensor.put(dst=routed_y_buf, peer=dst, ...)   # per-lane puts
    if dst != my_rank:
        pld.system.notify(target=combine_done, peer=dst, value=1, op=AtomicAdd)   # co-located notify
with pl.at(CORE_GROUP, deps=[push_tid], name_hint="combine_wait") as wait_tid:    # single-core wait
    for src ...: pld.system.wait(combine_done[src] >= moe_epoch*N_LOCAL, Ge)      # counting barrier
with pl.spmd(T//4, deps=[wait_tid], name_hint="combine_reduce"): ...

It compiles but fails at runtime: run_prepared failed code 507015, then AICore error 507018: bounded device drain failed, then a 300 s hang — i.e. inconsistent device-drain failures across runs. The same drain failure appears whether the notify is co-located in the spmd (above) or split into the barrier scope (my earlier attempt). Conclusion: the runtime/SDMA cannot drain many concurrent cross-rank TPUT/TNOTIFY channels issued from a multi-core pl.spmd scope (each lane gets its own channelGroupIdx = get_block_idx() channel; ~N_LOCAL cross-rank channels per rank, both ranks simultaneously). This is the capability that needs development: support multi-core (spmd) cross-rank pld.tensor.put + pld.system.notify that drains correctly (and/or a documented bound on concurrent cross-rank channels).

2. Minor codegen-validation gap — pto.comm.twait expected must be INT32, not INDEX

pld.system.wait(..., expected=moe_epoch * N_LOCAL) fails codegen with ptoas: custom op 'pto.comm.twait' invalid kind of type specified, because scalar * int promotes to INDEX (like offset arithmetic) and twait's expected operand requires INT32. Worked around with pl.cast(moe_epoch * N_LOCAL, pl.INT32). Low priority, but a clearer validation error ("twait expected must be INT32, got INDEX") — or accepting INDEX — would save debugging time.

Net

The achievable win is unchanged (~16-18 us/layer; combine_push is put-issue-bound, 71 cores idle). The ordering concern from the original post is resolved by the per-lane co-located put+notify counting-barrier design above — so the deliverable narrows to (1): make multi-core spmd cross-rank put/notify drain on device. Same limitation still blocks parallelizing the symmetric dispatch payload push.

[lwDavid]

Update 2 — twait works in spmd (once typed); the blocker is purely runtime delivery of multi-core cross-rank put/notify

Tested the all-in-one-spmd form (put + notify + wait all in a single pl.spmd scope), with the INT32 expected fix:

with pl.spmd(EP_WORLD_SIZE * N_LOCAL_EXPERTS, name_hint="combine_push") as push_tid:
    bucket = pl.tile.get_block_idx(); dst = bucket // N_LOCAL; e = bucket - dst*N_LOCAL
    ... for row in pl.range(n): pld.tensor.put(dst=routed_y_buf, peer=dst, ...)
    if dst != my_rank:
        pld.system.notify(target=combine_done, peer=dst, value=1, op=AtomicAdd)
    for src ...:                                   # wait now IN the same spmd scope
        pld.system.wait(combine_done[src] >= pl.cast(moe_epoch*N_LOCAL, pl.INT32), Ge)
with pl.spmd(T//4, deps=[push_tid], name_hint="combine_reduce"): ...

Findings:

• pto.comm.twait codegens fine inside a multi-core pl.spmd scope — the earlier compile failure was only the INDEX-vs-INT32 expected type bug (Update 1 item 2), not a scope restriction. So that codegen-validation note is the whole codegen story; nothing else blocks emitting put/notify/wait in spmd.
• Runtime still fails. All-in-spmd hangs (300s timeout): the cross-rank notifies issued from the spmd lanes never all land, so the counting-barrier waits never clear. This matches CP4's 507018 drain failure (separate single-core wait scope) — same root cause regardless of where the wait sits: the runtime/SDMA cannot deliver/drain the ~N_LOCAL concurrent cross-rank TPUT/TNOTIFY channels that a multi-core spmd scope issues (each lane = its own channelGroupIdx = get_block_idx()), with both ranks doing it simultaneously.

Net: the kernel/codegen side is fully solved (per-lane co-located put+notify+wait, counting barrier, INT32 threshold — all compile and are logically correct). The sole remaining work is the runtime capability: make N concurrent cross-rank pld.tensor.put + pld.system.notify issued from a multi-core pl.spmd scope deliver and drain correctly (or document the max concurrent cross-rank channel count and the required drain/quiesce). This is the single deliverable for both combine_push and the symmetric dispatch payload push.

[lwDavid]

Update 3 — minimal isolated repro: cross-rank notify+wait works in CORE_GROUP, fails in spmd (507015), even at 2 lanes

Reduced the whole thing to a ~90-line standalone barrier (no puts, no reduce, no data path): each rank does mutual cross-rank pld.system.notify(AtomicAdd) + pld.system.wait(Ge), once per lane, with a counting threshold. Same kernel, two scope variants, on a2a3 EP=2, with a passing single-core control:

variant	scope	result
control	pl.at(pl.Level.CORE_GROUP) (single core)	PASS
test	pl.spmd(2)	507015 bounded device drain failed
test	pl.spmd(32)	507015 bounded device drain failed

Key points:

• The only difference between PASS and FAIL is the scope: CORE_GROUP vs spmd. The barrier logic, counting threshold, windows, and harness are identical.
• It fails at lanes=2, i.e. the minimum multi-lane case — so this is not a "too many concurrent cross-rank channels" scaling issue. Issuing cross-rank notify/wait from any multi-core pl.spmd scope drain-fails.
• Deterministic (507015/507018 every run), fast-fail (not just the intermittent 12ms-style stall).

This is the clean root-cause isolation behind the earlier combine_push CP4/CP5 failures: it is the spmd scope + cross-rank system signalling combination the runtime cannot drain, independent of the puts. So the parallel-combine optimization needs the framework to support cross-rank notify/wait (and put) issued from a multi-core pl.spmd scope — the current codegen/runtime only supports them from a single-core InCore/CORE_GROUP scope.

Minimal repro available (pypto-lib models/deepseek/v4/min_barrier_draft.py, --mode single|spmd --lanes N); happy to port it into tests/st/distributed/ as an xfail if useful.

[lwDavid]

Root cause located and filed as a concrete codegen bug: #1913. It is not a missing capability — the orchestration codegen's Spmd path (GenerateSpmdCallCode) simply omits the per-DistributedTensor CommContext add_scalar(ext_<name>_ctx) that the InCore path emits (GenerateFunctionCallCode → EmitDistTensorCtxScalars, orchestration_codegen.cpp:2699), so the Spmd kernel reads a garbage CommContext at args[N] and the cross-rank notify/wait/put deadlocks (507015). Fixing #1913 should unblock the parallel combine_push / dispatch-push directly. Suggest tracking the work under #1913 and closing this once it lands.