perf(ep): rework how EP tuning picks and saves, remove small-token host overhead

[isytwu]

Summary

Small-token (4/8/16) EP16 dispatch/combine on MI300X is bound by the host
submission path, not by the kernels. The series started there and grew to
cover the tuning machinery that was supposed to measure it, which turned out to
be unable to pick a winner at all.

Runtime

1. perf(ep) — memoize the six output tensor views that dispatch() /
   combine() rebuilt on every call. This is the −32% below. Opt out with
   MORI_EP_DISABLE_VIEW_CACHE=1.

Tuning correctness

2. fix(ep) — the tuning margin was an absolute 1.0 GB/s. Where a whole
   sweep sits in the low single digits (4/8 tokens: 3.9 / 8.3 GB/s) nothing
   after the first candidate could ever clear it, so the sweep returned
   whatever it measured first; at 32 tokens (30.7 GB/s) it was under 4% and
   filtered almost nothing. Now relative, via MORI_EP_TUNING_MARGIN,
   defaulting to 0 (any improvement wins). The same bug and the same fix
   applied to the intra-node tuner, which shares the env var.
3. fix(ep) — selection used a slowest-rank bandwidth while the summary
   printed a grand mean, so a config could win the sweep on a number the table
   it was chosen from never showed. Both are now the grand-mean latency,
   the quantity printed as the table's "Average" row.
4. refactor(ep) — the keep-best save gate is gone. It compared
   new_bw > old_bw across two tuning runs from different hosts, ROCm
   versions and machine load, and on a "keep" it said so at logger.info,
   which this library's WARNING-level logger never emits. A discarded run was
   completely silent. Saves are now unconditional and always print what
   changed; the JSON's git diff is where the call actually gets made.
5. refactor(ep) — bandwidth_gbps is now the grand mean too, matching
   the printed table, and rules carry bandwidth_metric recording that. A
   re-tune that replaces a rule written under the old definition prints both
   numbers without a percentage instead of reporting a definition change as a
   20%+ win.
6. fix(ep) — MORI_TUNING_SCOPE=quick together with
   --save-tuning-config is now rejected, in both tuners and both batch
   scripts. quick sweeps a reduced grid (inter-node: 3 warp_per_block against
   full's 5; intra-node: 3 against 9, and powers of two only for block_num),
   so its winner is not a result worth committing. This changed the documented
   workflow — see Behaviour changes for existing users.
7. feat(ep) — tools/verify_tuned_config.sh, because the sweep only
   times candidates and never compares output: nothing confirmed the geometry
   it wrote was correct, only that it was fast.

Supporting

8. feat(ep) — MORI_EP_ROUNDS (default 10) now sets the rounds for both
   --cmd bench and each tuning candidate, so the two measure on equal
   footing. Round 0 is dropped as warmup; values below 2 are rejected.
9. feat(ep) — MORI_EP_TUNING_VERBOSE=1 prints each candidate's full
   table; block_num 8 and 16 added to the sweep's candidate list;
   rdma_worst_rank / ll_worst_rank printed next to the winner as a
   straggler check (not saved, not used for selection).
10. feat(ep) — tuned InterNodeV1LL EP16 configs for MI300X at hidden
    6144. Read the caveat in About the shipped configs before relying on
    these.
11. feat(ep) — test_low_latency.py's shape is configurable instead of
    hardcoded.
12. Tooling — docker exec is probed rather than assuming sudo;
    --local-docker runs rank 0 in a container so the driver can stay on the
    host and its SSH keys never enter a container; --rdma-tc plumbs
    MORI_RDMA_TC through to both ranks (it was silently dropped —
    build_torchrun_cmd names the environment explicitly and the ranks start
    via docker exec/ssh, neither of which inherits the caller's); the
    cleanup pkill chain no longer kills itself and is scoped to this
    script rather than every torchrun the user owns on the node; run_all_internode_tuning.sh
    propagates a non-zero exit; /logs/ is gitignored.

No kernel changes — see What is not here.

Measurements

EP16 = 2 × 8 MI300X, v1_ll, hidden 6144, fp8_e4m3_fnuz dispatch / bf16
combine, topk 8, --num-qp 1. Medians over 3–5 runs; single runs are not
meaningful here (see Measurement caveats).

Final validation with test_low_latency.py --num-tokens 4 --hidden 6144:

	pre-change	this PR
end-to-end bench() avg_t	216.8 us	147.4 us	−32%
end-to-end min_t	199.5 us	134.2 us	−33%
bench_kineto dispatch	62.0 us	59.7 us	unchanged
bench_kineto combine	100.2 us	96.3 us	unchanged

The last two rows are the point: GPU kernel time does not move, the entire
gain is host-side. bench_kineto explicitly uses a large kernel plus a barrier
to "eliminate the unbalanced CPU launch overhead", so it is structurally
blind to this class of problem — which is why it went unnoticed.

Same picture from test_dispatch_combine_internode.py, A/B in a single
session:

	pre-change	this PR
combine	144.7 us	86.8 us	−40%
wall/round	220.9 us	162.4 us	−26%
host enqueue/round	204.3 us	142.7 us	−30%
dispatch	81.8 us	74.2 us	−9%

Reproduced on a second, independent node pair with the same numbers (combine
150→88, wall 222→164, enqueue 211→146), so this is not a property of one
machine.

Why the regime is overhead-bound

test	result	conclusion
hidden 1024 vs 6144 (6× data)	93→85 us / 147→145 us — unchanged	not data-bound
grid 304 → 24 blocks	12.13 → 12.05 us	not grid-bound
host enqueue vs wall	204 us vs 221 us per round	host-bound

Latency is also flat across 4/8/16 tokens.

Behaviour changes for existing users

• --tuning-scope quick can no longer write a config. The documented
  "quick sweep" workflow in MORI-EP-GUIDE.md did exactly that, and
  run_all_internode_tuning.sh defaulted to quick while always passing a
  config output — so with no arguments it would now have failed all 42 jobs.
  Its default is therefore flipped to full, the guide is updated, and both
  batch scripts reject the combination up front (in a second, rather than after
  a two-node launch) naming the two knobs that resolve it. quick remains
  available for exploration with --config-output ''.
• Saving overwrites unconditionally. A re-tune that measures worse will now
  replace a good rule. That is deliberate — the alternative failed silently —
  but it means the printed delta is the thing to read, and a bad run should be
  discarded with git checkout rather than trusted to be filtered out.
• Tuning rounds default 5 → 10 (MORI_EP_ROUNDS), so a sweep takes about
  twice as long. That buys a tuner that can distinguish configs at all.

About the shipped configs

The configs are a reproducible starting point, not a proven optimum. The
tuner measures in the same host-bound regime this PR is about, so at 4 tokens
its metric cannot cleanly separate geometries.

Re-tuning the same four shapes three times on one node pair reproduced 6 of the
8 saved geometries exactly and put the other two in a flat region (latency
within 1-4% of the saved rule). Two of the 24 measurements were whole-sweep
degradations of 25-31% that landed on a different shape each run -- in both
cases the fastest of all 105 candidates for that shape was the degraded one,
so the tuner picked correctly and the run itself was bad. The mandatory save
report is what made those visible; under the old gate they would have been
discarded in silence.

verify_tuned_config.sh reports, before running anything, which phases have a
rule that will actually be selected, and refuses to claim it verified a phase
that fell back to the built-in geometry:

  dispatch rules:  4 (dtypes: fp8_e4m3_fnuz) queried with fp8_e4m3_fnuz -> MATCH
  combine  rules:  4 (dtypes: bf16) queried with bf16 -> MATCH

Review notes on the memoization

The addresses come from _dispatch_out_ptrs / _combine_out_ptrs, which
__init__ already snapshots once from handle-owned symmetric shmem
(GetShmemDispatchOutTokMemObj() etc., selected by the immutable
config.kernelType, allocated at handle construction with no realloc path).
The pointer-stability assumption is pre-existing — this change caches a
view over pointers the code already treats as fixed, and does not add a
reallocation guarantee (a realloc would leave both the key and the view stale,
exactly as before).

• Growth: bounded by the distinct (shape, dtype) an op is called with —
  one, for a fixed model.
• Lifetime: entries alias handle-owned memory and live on the op, so
  nothing is retained beyond it.
• Object identity: callers now get the same object rather than an equal
  one. Both alias the same buffer and the kernels overwrite it in place, so the
  only visible change is identity. Any caller that expected two dispatch()
  results to be independent snapshots was already broken.
• Concurrency: worst case two threads build the same view and one wins;
  both are equivalent.
• Opt-out: MORI_EP_DISABLE_VIEW_CACHE=1, read once in __init__ —
  probing the environment per call would reintroduce the cost this removes.

Correctness: --cmd test at 4 and 16 tokens, 500 rounds, on all 16 ranks,
completed without raising. Note that the error times: N line those runs print
is not the evidence — error_round is never populated anywhere in the file,
so that counter is always 0. Every comparison in run_test_once is a bare
assert, so the signal is that the run finished at all.

--num-qp: 1 is best, and here is why

Swept with the host off the critical path so the parameter can actually show,
5 runs each:

--num-qp	1	2	3	4
round (us)	149.8	160.1	168.4	170.8

Monotonic, and the distributions for 1 and 4 do not overlap (qp1 max 162.2 <
qp4 min 165.6). It holds at every size tested: −12.3% at 4 tokens, −3.1% at
128, −5.1% at 512.

The mechanism is in the combine cross-device barrier, which issues one RDMA
atomic per QP:

for (int i = 0; i < config.numQpPerPe; i++) {
  shmem::ShmemAtomicTypeNonFetchThread<uint64_t>(..., core::AMO_ADD, proxyPe, i);
}

Meanwhile the data path picks a QP with
qpId = (tokenId / warpSize) % numQpPerPe, which is always 0 at 4 tokens. So
extra QPs buy nothing and cost extra RDMA atomic issues — span-profiling puts a
single GPU-initiated RDMA issue at ~13–16 us, which matches the ~21 us gap
between qp=1 and qp=4.

Two notes: measured plainly, this is invisible (158.4 vs 161.8 us, inside the
noise) because the host dominates. And MORI_NUM_QP_PER_PE (transport-level QP
count, default 4) is a different knob and had no measurable effect (152.7 vs
151.3 us); only the op config's numQpPerPe matters.

This is a usage recommendation; the default is left alone here.

Reproducing

Two nodes, $R = 0 / 1, start node 1 first.

export MORI_RDMA_TC=41 GPU_PER_NODE=8
export GLOO_SOCKET_IFNAME=<IFNAME> MORI_SOCKET_IFNAME=<IFNAME>

# baseline: pre-change behaviour
export MORI_EP_DISABLE_VIEW_CACHE=1
torchrun --nnodes=2 --node_rank=$R --nproc_per_node=1 \
  --master_addr=<HOST0> --master_port=1234 \
  examples/ops/dispatch_combine/test_dispatch_combine_internode.py \
  --cmd bench --kernel-type v1_ll --num-qp 1 \
  --max-tokens 4 --hidden-dim 6144 \
  --dtype fp8_e4m3_fnuz --combine-dtype bf16 --quant-type none \
  --block-num 32 --warp-per-block 4 --rdma-block-num 8

# this PR: drop the opt-out and rerun the identical torchrun command
unset MORI_EP_DISABLE_VIEW_CACHE

Tuning and verification, driver on the host with mori in containers on both
nodes:

bash tools/batch_internode_tuning.sh \
    --master-addr <HOST0> --peer-host <USER>@<HOST1> --ifname <IFNAME> \
    --local-docker <RANK0_CONTAINER> --docker <PEER_CONTAINER> --rdma-tc 41 \
    --kernel-type v1_ll --num-qp 1 \
    --dtype fp8_e4m3_fnuz --combine-dtype bf16 --quant-type none \
    --tokens-list "4,8,16,32" --hidden-dims "6144"

bash tools/verify_tuned_config.sh <same connection and shape flags>

End-to-end numbers above come from test_low_latency.py --num-tokens 4 --hidden 6144, which this series makes configurable.

Measurement caveats, all of which produced wrong conclusions during this
work:

• The first run after a JIT-cache wipe or an idle period reads 40–80% slow.
  Discard one run before comparing.
• Run-to-run spread is 10–20%, and roughly 1 run in 5 lands in a reproducible
  slow mode (~1.5×). Compare medians over ≥3 runs, never a single one.
• EpDispatchCopyToStaging makes a good canary: nothing here touches it, so a
  reading above ~5 us means the environment is bad rather than the patch.

What is not here

No kernel optimization. Four kernel changes were implemented, verified
correct, and measured against per-kernel GPU time from the CUDA profiler — all
within noise (136.0/136.2 vs 135.5/135.5 us): a MultiWarpIter minimum slice,
skipping combine padding slots, an s_sleep poll backoff, and counting only
active blocks in the dispatch send barrier.

Where the GPU time actually goes, at 4 tokens / hidden 6144 (136 us total): the
two LL kernels are 82% of it, and fitting the token sweep gives dispatch ≈
50 us + 0.5 us/token and combine ≈ 46 us + 0.48 us/token. At 4 tokens
that intercept is 98% of the cost.

That fixed cost is not the fabric — ib_write_lat on the same NIC pair is 10.7
us for the 24 KB a 4-token dispatch moves. Span-profiling with
MORI_TRACE_SPAN locates it:

• one warp spends 16.2 us issuing a single RDMA put while the other 127
  leave the send phase in 1.24 us; the recv phase's 17.25 us warp-ramp is
  exactly that warp arriving late
• combine's median warp does 14.26 us of real work — warpsPerToken is
  pinned at 4, so each warp gathers 8 experts × 1536 elements over XGMI

The put cannot simply be trimmed: downgrading its four __threadfence_system()
calls to agent scope regresses everything (dispatch 55 → 131 us) because the
NIC stops seeing the payload.

Worth filing separately:

1. ENABLE_PROFILER=ON does not build on main — intranode_ll.hpp
   references three Slot enumerators that are never declared, because the
   generator groups slots by source-file stem while Slot there aliases
   intranode.hpp's enum.
2. Tuning configs frequently do not take effect, three independent reasons:
   MORI_EP_LAUNCH_CONFIG_MODE defaults to MANUAL, so the shipped JSONs are
   inert unless a caller opts in; tuning writes to the repo while the op reads
   the installed copy; on a miss _find_fallback_config silently substitutes
   another GPU model's file (observed running mi308x configs on MI300X),
   announcing it at logger.info. Each of these is a silent fallback with no
   warning a user will see.
3. get_launch_config looks both phases up with config.data_type, so for a
   mixed-dtype config it returns nothing for combine. It is not on the runtime
   path -- dispatch()/combine() go through _resolve_launch_params with
   dtype=input.dtype, which is correct -- and nothing in the repository calls
   it, so this is latent. Documented in place rather than changed, since fixing
   it means deciding what its signature should be.
4. error_round in test_dispatch_combine_internode.py is passed to
   run_test_once and never written, so every error times: line the test
   prints reads 0 regardless of outcome. It reads like a correctness summary
   and is not one.