From 32536daf2e758f5bea75b84428c8d01264436d53 Mon Sep 17 00:00:00 2001 From: Yangwook Kang Date: Tue, 14 Apr 2026 00:38:41 -0700 Subject: [PATCH] Fix ADR-0025: IPCQ direction addressing via address-based matching MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 2-rank bidirectional ring deadlock: when E and W neighbors point to the same peer, sender-coord matching in _handle_meta_arrival / _credit_worker picked the first direction in dict order, landing data in the wrong rx slot relative to what the kernel recv(W) was waiting on. Fix (ADR-0025 D1/D2/D3): - install.reverse_direction: prefer OPPOSITE direction (E↔W, N↔S) when peer has it pointing back to us; fallback to any matching for topologies without opposite convention (tree_binary parent/child). - _handle_meta_arrival: match by token.dst_addr range against each qp's my_rx_base_pa + n_slots × slot_size window (unambiguous). - _credit_worker: match by credit.dst_rx_base_pa == qp.peer.rx_base_pa. - IpcqCreditMetadata: new dst_rx_base_pa field carrying receiver-side rx base; _delayed_credit_send fills it from the consuming qp. Tests (Phase 1 → Phase 2): - test_reverse_direction_opposite_preference_2rank_ring - test_reverse_direction_opposite_preference_4rank_ring_sanity - test_meta_arrival_matches_by_dst_addr_same_peer - test_credit_matches_by_dst_rx_base_pa_same_peer - Existing credit-return test updated with dst_rx_base_pa. 508 tests pass. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/kernbench/ccl/install.py | 23 ++- src/kernbench/common/ipcq_types.py | 9 +- src/kernbench/components/builtin/pe_ipcq.py | 38 +++-- tests/test_ccl_install.py | 60 ++++++++ tests/test_ipcq_types.py | 6 +- tests/test_pe_ipcq.py | 158 +++++++++++++++++++- 6 files changed, 277 insertions(+), 17 deletions(-) diff --git a/src/kernbench/ccl/install.py b/src/kernbench/ccl/install.py index 4218763..6607048 100644 --- a/src/kernbench/ccl/install.py +++ b/src/kernbench/ccl/install.py @@ -219,9 +219,24 @@ def install_ipcq( "neighbor_table": neighbor_table, } - def reverse_direction(my_rank: int, peer_rank: int) -> str | None: - """Find which direction in peer's neighbor table points back to my_rank.""" - for d, target in neighbor_table[peer_rank].items(): + _OPPOSITE_DIR = {"E": "W", "W": "E", "N": "S", "S": "N"} + + def reverse_direction(my_rank: int, peer_rank: int, my_dir: str) -> str | None: + """Find peer's direction that reciprocates my_dir→peer_rank. + + Prefer the OPPOSITE direction (E↔W, N↔S) when the peer has it + pointing back to us (ADR-0025 D1). This matters in 2-rank + bidirectional rings where both E and W on one side point to the + same peer — without the preference, dict-order first-match would + route data into the wrong rx slot. Falls back to any direction + pointing back for topologies without an opposite convention + (e.g. tree_binary's parent/child). + """ + nt = neighbor_table[peer_rank] + opp = _OPPOSITE_DIR.get(my_dir) + if opp is not None and nt.get(opp) == my_rank: + return opp + for d, target in nt.items(): if target == my_rank: return d return None @@ -234,7 +249,7 @@ def install_ipcq( if peer_rank is None: continue peer_s, peer_c, peer_p = rank_pe[peer_rank] - peer_dir = reverse_direction(r, peer_rank) + peer_dir = reverse_direction(r, peer_rank, d) if peer_dir is None: # Peer doesn't have a reverse entry — skip (asymmetric topology) continue diff --git a/src/kernbench/common/ipcq_types.py b/src/kernbench/common/ipcq_types.py index 2cddead..578aaed 100644 --- a/src/kernbench/common/ipcq_types.py +++ b/src/kernbench/common/ipcq_types.py @@ -196,10 +196,17 @@ class IpcqCreditMetadata: Sent by ``PeIpcqComponent._delayed_credit_send`` after a bottleneck-BW based latency, putting the metadata directly into the peer's pre-wired credit store (no fabric routing). + + ``dst_rx_base_pa`` is the receiver's ``my_rx_base_pa`` for the direction + whose slot was consumed. The original sender matches this against + ``qp.peer.rx_base_pa`` to find the correct direction (ADR-0025 D3) — + unambiguous even when multiple directions share the same peer (e.g. + 2-rank bidirectional ring). """ consumer_seq: int # my_tail at recv side (new tail value) - src_sip: int # which peer is sending the credit + dst_rx_base_pa: int # receiver-side my_rx_base_pa (ADR-0025 D3) + src_sip: int # which peer is sending the credit (diag) src_cube: int src_pe: int src_direction: str # sender-side direction (peer maps to its own) diff --git a/src/kernbench/components/builtin/pe_ipcq.py b/src/kernbench/components/builtin/pe_ipcq.py index 27b5d8c..5e31876 100644 --- a/src/kernbench/components/builtin/pe_ipcq.py +++ b/src/kernbench/components/builtin/pe_ipcq.py @@ -370,11 +370,21 @@ class PeIpcqComponent(ComponentBase): # ── Metadata arrival from PE_DMA (D9) ── def _handle_meta_arrival(self, msg: IpcqMetaArrival) -> None: + """Match arrival to the correct direction by dst_addr range (ADR-0025 D2). + + Each direction has a unique rx buffer address range + ([my_rx_base_pa, my_rx_base_pa + n_slots * slot_size)). The token's + dst_addr (set by the sender's IPCQ when computing the peer slot + address) falls within exactly one such range. Address-based matching + is unambiguous even when multiple directions share the same peer + (2-rank bidirectional ring). + """ token = msg.token - sender_key = (token.src_sip, token.src_cube, token.src_pe) + dst_addr = token.dst_addr for d, qp in self._queue_pairs.items(): - p = qp["peer"] - if (p.sip, p.cube, p.pe) == sender_key: + base = qp["my_rx_base_pa"] + size = qp["n_slots"] * qp["slot_size"] + if base <= dst_addr < base + size: qp["peer_head_cache"] = max(qp["peer_head_cache"], token.sender_seq + 1) # Track arrived token for strict-mode peek self._arrived_tokens.setdefault(d, []).append(token) @@ -391,19 +401,22 @@ class PeIpcqComponent(ComponentBase): if not ev.triggered: ev.succeed() return - # Unknown sender — silently drop (could log) + # Unknown dst_addr — silently drop (could log) # ── Credit return (fast path) ── def _credit_worker(self, env: simpy.Environment) -> Generator: - """Process IpcqCreditMetadata from credit_inbox.""" + """Process IpcqCreditMetadata from credit_inbox. + + Matches credit to the correct direction by `credit.dst_rx_base_pa == + qp.peer.rx_base_pa` (ADR-0025 D3). This is unambiguous even when + multiple directions share the same peer (2-rank bidirectional ring). + """ assert self._credit_inbox is not None while True: credit: IpcqCreditMetadata = yield self._credit_inbox.get() - sender_key = (credit.src_sip, credit.src_cube, credit.src_pe) for d, qp in self._queue_pairs.items(): - p = qp["peer"] - if (p.sip, p.cube, p.pe) == sender_key: + if qp["peer"].rx_base_pa == credit.dst_rx_base_pa: qp["peer_tail_cache"] = max(qp["peer_tail_cache"], credit.consumer_seq) # Wake any blocked send on this direction waiters = self._send_waiters.get(d, []) @@ -421,12 +434,19 @@ class PeIpcqComponent(ComponentBase): new_tail: int, ) -> Generator: """Wait bottleneck-BW latency, then put IpcqCreditMetadata into peer - credit store (D9 fast path).""" + credit store (D9 fast path). + + Carries ``dst_rx_base_pa`` = this PE's my_rx_base_pa for the + consumed direction. The peer (original sender) matches this against + qp.peer.rx_base_pa to identify the correct qp (ADR-0025 D3). + """ latency_ns = self._credit_latency_ns(direction) if latency_ns > 0: yield env.timeout(latency_ns) + qp = self._queue_pairs[direction] meta = IpcqCreditMetadata( consumer_seq=new_tail, + dst_rx_base_pa=qp["my_rx_base_pa"], src_sip=self._self_sip, src_cube=self._self_cube, src_pe=self._self_pe, diff --git a/tests/test_ccl_install.py b/tests/test_ccl_install.py index 738611e..f03e6b5 100644 --- a/tests/test_ccl_install.py +++ b/tests/test_ccl_install.py @@ -98,3 +98,63 @@ def test_install_ipcq_credit_stores_wired(): qp_e = pe0.queue_pairs["E"] assert qp_e["peer_credit_store"] is pe1.credit_inbox + + +# ── ADR-0025 D1: reverse_direction opposite-preference ─────────────── + + +def test_reverse_direction_opposite_preference_2rank_ring(): + """ADR-0025 D1: In a 2-rank bidirectional ring both E and W point to the + same peer; reverse_direction must pick the OPPOSITE direction (W for E, + E for W) so rx_base targets the semantically-correct slot. + + Concretely: rank 0 sending via E to rank 1 must target rank 1's W-rx + buffer (not rank 1's E-rx), because rank 1's kernel recv(W) reads from + its W-rx. + """ + engine, topo = _engine() + cfg = load_ccl_config() + merged = resolve_algorithm_config(cfg, name="ring_allreduce_tcm") + merged["world_size"] = 2 + install_ipcq(engine, topo.spec, merged) + + ipcq0 = engine._components["sip0.cube0.pe0.pe_ipcq"] + ipcq1 = engine._components["sip0.cube0.pe1.pe_ipcq"] + + rank1_e_rx = ipcq1.queue_pairs["E"]["my_rx_base_pa"] + rank1_w_rx = ipcq1.queue_pairs["W"]["my_rx_base_pa"] + + qp0_e = ipcq0.queue_pairs["E"] + qp0_w = ipcq0.queue_pairs["W"] + + # rank 0's E entry should target rank 1's W-rx (opposite), NOT rank 1's E-rx. + assert qp0_e["peer"].rx_base_pa == rank1_w_rx, ( + f"expected rank 0's E peer.rx_base_pa == rank 1's W-rx ({rank1_w_rx:#x}), " + f"got {qp0_e['peer'].rx_base_pa:#x} (matches E-rx: {rank1_e_rx:#x}) — " + f"reverse_direction picked same-label instead of opposite" + ) + # rank 0's W entry should target rank 1's E-rx (opposite). + assert qp0_w["peer"].rx_base_pa == rank1_e_rx + + +def test_reverse_direction_opposite_preference_4rank_ring_sanity(): + """ADR-0025 D1 sanity: ws>=3 ring. E and W have distinct peers, so + opposite-preference produces same result as old dict-order first-match. + This test should PASS both under current and post-fix code. + """ + engine, topo = _engine() + cfg = load_ccl_config() + merged = resolve_algorithm_config(cfg, name="ring_allreduce_tcm") + merged["world_size"] = 4 + install_ipcq(engine, topo.spec, merged) + + ipcq0 = engine._components["sip0.cube0.pe0.pe_ipcq"] + ipcq1 = engine._components["sip0.cube0.pe1.pe_ipcq"] + ipcq3 = engine._components["sip0.cube0.pe3.pe_ipcq"] + + # rank 0 E → rank 1 → rank 1's W-rx + qp0_e = ipcq0.queue_pairs["E"] + assert qp0_e["peer"].rx_base_pa == ipcq1.queue_pairs["W"]["my_rx_base_pa"] + # rank 0 W → rank 3 (last in ring) → rank 3's E-rx + qp0_w = ipcq0.queue_pairs["W"] + assert qp0_w["peer"].rx_base_pa == ipcq3.queue_pairs["E"]["my_rx_base_pa"] diff --git a/tests/test_ipcq_types.py b/tests/test_ipcq_types.py index 647837b..5810db4 100644 --- a/tests/test_ipcq_types.py +++ b/tests/test_ipcq_types.py @@ -63,7 +63,8 @@ def test_ipcq_dma_token(): def test_ipcq_credit_metadata(): cm = IpcqCreditMetadata( - consumer_seq=3, src_sip=0, src_cube=0, src_pe=1, src_direction="W", + consumer_seq=3, dst_rx_base_pa=0x1000, + src_sip=0, src_cube=0, src_pe=1, src_direction="W", ) assert cm.consumer_seq == 3 assert cm.src_direction == "W" @@ -71,7 +72,8 @@ def test_ipcq_credit_metadata(): def test_ipcq_credit_metadata_frozen(): cm = IpcqCreditMetadata( - consumer_seq=3, src_sip=0, src_cube=0, src_pe=1, src_direction="W", + consumer_seq=3, dst_rx_base_pa=0x1000, + src_sip=0, src_cube=0, src_pe=1, src_direction="W", ) with pytest.raises(Exception): cm.consumer_seq = 99 # type: ignore diff --git a/tests/test_pe_ipcq.py b/tests/test_pe_ipcq.py index b339315..91d38c1 100644 --- a/tests/test_pe_ipcq.py +++ b/tests/test_pe_ipcq.py @@ -291,9 +291,12 @@ def test_send_blocks_when_peer_slot_full(): env.run(until=20) assert not req5.done.triggered - # Send a credit return: peer (E direction, pe=1) consumed slot 0 + # Send a credit return: peer (E direction, pe=1) consumed slot 0. + # dst_rx_base_pa is the peer-side rx buffer — which equals my qp_E's + # peer.rx_base_pa (0x10_000 from _install_two_neighbors). credit = IpcqCreditMetadata( consumer_seq=1, # peer consumed up to my_tail=1 + dst_rx_base_pa=0x10_000, # E's peer.rx_base_pa (ADR-0025 D3) src_sip=0, src_cube=0, src_pe=1, src_direction="W", # peer's view ) comp.credit_inbox.put(credit) @@ -315,3 +318,156 @@ def test_init_installs_neighbors(): assert comp._queue_pairs["W"]["peer"].pe == 2 assert comp._queue_pairs["E"]["my_head"] == 0 assert comp._queue_pairs["E"]["peer_tail_cache"] == 0 + + +# ── ADR-0025: address-based matching in meta arrival / credit ──────── + + +def _install_same_peer_neighbors( + env: simpy.Environment, comp: PeIpcqComponent, +) -> tuple[simpy.Store, simpy.Store]: + """Install E and W neighbors BOTH pointing to the same peer (pe=1). + + This mirrors the 2-rank bidirectional ring topology (ADR-0025 motivation): + rank 0's E and W neighbors are the same peer rank, but target different + rx slots on that peer (E→peer's W-rx, W→peer's E-rx). + + - E's peer.rx_base_pa = 0x10_000 (peer's W-rx buffer) + - W's peer.rx_base_pa = 0x20_000 (peer's E-rx buffer) + - my_rx_base_pa: E=0x30_000, W=0x40_000 (local rx for each dir) + """ + peer_e_credit = simpy.Store(env) + peer_w_credit = simpy.Store(env) + + ep_e = IpcqEndpoint( + sip=0, cube=0, pe=1, + buffer_kind="tcm", + rx_base_pa=0x10_000, rx_base_va=0, + n_slots=4, slot_size=4096, + ) + ep_w = IpcqEndpoint( + sip=0, cube=0, pe=1, # SAME peer as ep_e + buffer_kind="tcm", + rx_base_pa=0x20_000, rx_base_va=0, # different target slot + n_slots=4, slot_size=4096, + ) + init_msg = IpcqInitMsg( + correlation_id="t", request_id="t", + target_sips=(0,), target_cubes=(0,), target_pe=0, + entries=( + IpcqInitEntry( + direction="E", peer=ep_e, + my_rx_base_pa=0x30_000, my_rx_base_va=0, + n_slots=4, slot_size=4096, + peer_credit_store=peer_e_credit, + ), + IpcqInitEntry( + direction="W", peer=ep_w, + my_rx_base_pa=0x40_000, my_rx_base_va=0, + n_slots=4, slot_size=4096, + peer_credit_store=peer_w_credit, + ), + ), + backpressure_mode="sleep", + buffer_kind="tcm", + credit_size_bytes=16, + ) + done = env.event() + comp.in_ports["host"].put(_FakeTxn(request=init_msg, done=done)) + env.run(until=done) + return peer_e_credit, peer_w_credit + + +def test_meta_arrival_matches_by_dst_addr_same_peer(): + """ADR-0025 D2: when E and W point to the same peer (2-rank ring), + dst_addr range must determine which qp's peer_head_cache updates. + + Under the old sender-key matching, the first matching direction (E) + would win for any arrival, regardless of which rx slot was written. + Under D2 address-based matching, dst_addr within W's rx range + (my_rx_base_pa_W .. +n_slots*slot_size) must update W, and dst_addr + within E's rx range must update E. + """ + env = simpy.Environment() + comp = _make_pe_ipcq(env) + _install_same_peer_neighbors(env, comp) + + # Arrival into W's rx buffer (my_rx_base_pa=0x40_000) + token_into_w = IpcqDmaToken( + src_addr=0, src_space="tcm", + dst_addr=0x40_000, dst_endpoint=comp._queue_pairs["W"]["peer"], + nbytes=64, handle_id="w1", + shape=(8,), dtype="f16", + sender_seq=0, + src_sip=0, src_cube=0, src_pe=1, src_direction="E", + ) + comp.in_ports["host"].put(IpcqMetaArrival(token=token_into_w)) + env.run(until=5) + + # W's peer_head_cache should increment; E's stays 0. + assert comp._queue_pairs["W"]["peer_head_cache"] == 1, ( + "W qp should have been updated because dst_addr is in W's rx range" + ) + assert comp._queue_pairs["E"]["peer_head_cache"] == 0, ( + "E qp should NOT be updated; current sender-key matching wrongly " + "picks the first direction with a matching peer" + ) + + # Second arrival into E's rx buffer (my_rx_base_pa=0x30_000) + token_into_e = IpcqDmaToken( + src_addr=0, src_space="tcm", + dst_addr=0x30_000, dst_endpoint=comp._queue_pairs["E"]["peer"], + nbytes=64, handle_id="e1", + shape=(8,), dtype="f16", + sender_seq=0, + src_sip=0, src_cube=0, src_pe=1, src_direction="W", + ) + comp.in_ports["host"].put(IpcqMetaArrival(token=token_into_e)) + env.run(until=10) + + assert comp._queue_pairs["E"]["peer_head_cache"] == 1 + assert comp._queue_pairs["W"]["peer_head_cache"] == 1 + + +def test_credit_matches_by_dst_rx_base_pa_same_peer(): + """ADR-0025 D3: credit must carry dst_rx_base_pa (the receiver-side + rx buffer base) so the original sender can match it against + qp.peer.rx_base_pa and find the correct direction. Under old + sender-key matching, first-match-wins would always pick E when + E and W share the same peer. + """ + env = simpy.Environment() + comp = _make_pe_ipcq(env) + _install_same_peer_neighbors(env, comp) + + # Credit corresponding to a send through W direction: + # - My W sent to peer's rx at 0x20_000 (qp_w["peer"].rx_base_pa) + # - Peer consumed it; sends credit back with dst_rx_base_pa=0x20_000 + # - Receiver (me, the original sender) should update W's peer_tail_cache + credit_for_w = IpcqCreditMetadata( + consumer_seq=1, + dst_rx_base_pa=0x20_000, # matches W's peer.rx_base_pa + src_sip=0, src_cube=0, src_pe=1, src_direction="E", + ) + comp.credit_inbox.put(credit_for_w) + env.run(until=5) + + assert comp._queue_pairs["W"]["peer_tail_cache"] == 1, ( + "W's peer_tail_cache should update — credit.dst_rx_base_pa matches " + "W qp's peer.rx_base_pa" + ) + assert comp._queue_pairs["E"]["peer_tail_cache"] == 0, ( + "E's peer_tail_cache should NOT update" + ) + + # Second credit: for E direction + credit_for_e = IpcqCreditMetadata( + consumer_seq=2, + dst_rx_base_pa=0x10_000, # matches E's peer.rx_base_pa + src_sip=0, src_cube=0, src_pe=1, src_direction="W", + ) + comp.credit_inbox.put(credit_for_e) + env.run(until=10) + + assert comp._queue_pairs["E"]["peer_tail_cache"] == 2 + assert comp._queue_pairs["W"]["peer_tail_cache"] == 1