fix(ep/intranode training): restore replay-mode payload dispatch dropped by #532

[sudhu2k]

Motivation

Restore correctness of replay-mode dispatch, which is used in MoE training
backward pass. Since #532 (ae9635f7, "feat(ep/intranode): bring MoE dispatch/combine
up on gfx1250 (MI450)"), intra-node replay-mode dispatch moved no payload, so the
scattered activations/gradients were left stale. In Megatron-LM MoE training this surfaced
as an exploding gradient norm and immediate divergence,
while the forward loss stayed exactly correct.

The forward pass was unaffected because it uses cache-routing dispatch + replay-routing
combine (both intact). Only the backward pass uses replay-routing dispatch
(MoriCombine.backward → op.dispatch(routing=...)), which was the single broken path.

Technical Details

#532 refactored EpDispatchIntraNodeKernel_body (src/ops/dispatch_combine/intranode.hpp)
into three block-cooperative phases (count tokens → reserve slots → scatter payload), and
gated all three on if (args.tokenIndices && args.inpTokenBuf && !args.replayMode).
This inadvertently dropped the pre-#532 else branch that handled replay routing, so in
replay mode the kernel skipped every phase and moved zero payload into dispatchOut
(grad_x stayed stale).

Fix: re-add the replay branch. In replay mode we recover (destPe, destTokId) from the
dispDestTokIdMap written by the matching cache-routing dispatch, then copy metadata
(indices, weights, scales) and the token payload — exactly the pre-#532 behavior.
destPeTokenCounter / dispTokOffset are intentionally left untouched: in replay mode the
host takes total_recv from the routing handle rather than the kernel's signal, and the
completion handshake is a pure ordering barrier.

• Scope: intra-node kernel only.

Why existing tests missed it

test_replay_correctness exercised replay-dispatch but only asserted routing-map
immutability and the recv token count. A replay dispatch that moves zero payload leaves
the routing maps untouched and preserves the count, so all assertions passed while no
activation was scattered.

Regression guard added: after the replay dispatch, route the same payload through a
fresh op on the fully independent default path (no routing handle) and compare downstream
combine outputs. Combine outputs are in source-token order, so the check is invariant to
the recv-slot assignment that differs across ops; a broken replay (stale/zero payload)
diverges by ~activation magnitude and trips the assertion.

Files changed

• src/ops/dispatch_combine/intranode.hpp — restore replay-mode payload dispatch.
• tests/python/ops/test_dispatch_combine_routing_handle.py — payload guard in
  _replay_correctness (covers IntraNode + InterNodeV1, single- and two-node).

Test Plan

• Unit: pytest tests/python/ops/test_dispatch_combine_routing_handle.py -k replay_correctness
  on 8 GPUs, covering:
  • test_replay_correctness[intra] — IntraNode.
  • test_replay_correctness[v1] — InterNodeV1, single-node.
  • test_replay_correctness_v1_two_nodes — InterNodeV1, 2-node RDMA (gpu_per_node=4).

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.