# Copyright (c) 2026 LightSeek Foundation # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. """EPD encode->prefill transfer: wire-codec round-trips, row-shard geometry, and prefill receive-buffer sizing. Pure-logic silent-corruption guards (the RDMA write itself is unchecked); no GPU, no live Mooncake hop.""" from __future__ import annotations import threading import pytest import torch from tokenspeed.runtime.multimodal.inputs import ( Modality, MultimodalDataItem, ) from tokenspeed.runtime.pd.base.status import TransferPoll from tokenspeed.runtime.pd.epd.embedding_transfer import ( MooncakeEmbeddingSender, shard_payload, validate_fanout_frames, ) from tokenspeed.runtime.pd.epd.entities import ( REGISTER_ROOM_SENTINEL, EmbeddingArgsRegisterInfo, EmbeddingChunk, EmbeddingTransferError, EmbeddingTransferInfo, ) from tokenspeed.runtime.pd.epd.prefill_receiver import ( DONE, receive_encoded_embeddings, shard_rows, start_embedding_receive, ) # Each merged concern keeps its own per-test setup (the shard # pointer-math tests need the TOKENSPEED_EPD_RECV_POOL_SLOTS=0 path, the # receive tests need a small pool); route by test name. _RECV_TESTS = { "test_receive_sizes_buffers_per_item_no_deepstack", "test_receive_allocates_deepstack_columns", "test_receive_skips_already_encoded_item_on_recall", } @pytest.fixture(autouse=True) def _epd_transfer_env(request, monkeypatch): if request.node.name in _RECV_TESTS: _recv_setup(monkeypatch) else: _shard_setup(monkeypatch) # A high device VA to catch any 32-bit truncation / wrong struct format. HIGH_PTR = 0xFFFF_FFFF_FFFF_F000 MID_PTR = 0x7FFF_FFFF_0000 def test_register_info_round_trip(): info = EmbeddingArgsRegisterInfo( room=REGISTER_ROOM_SENTINEL, endpoint="10.0.0.7", dst_port=5123, mooncake_session_id="10.0.0.7:41999", dst_embedding_ptr=HIGH_PTR, dst_deepstack_ptr=MID_PTR, ) frames = info.to_zmq() assert len(frames) == 6 # layout lock assert frames[0] == b"None" assert EmbeddingArgsRegisterInfo.from_zmq(frames) == info def test_transfer_info_round_trip_with_deepstack(): info = EmbeddingTransferInfo( room=42, endpoint="10.0.0.7", dst_port=5123, mooncake_session_id="10.0.0.7:41999", dst_embedding_ptr=HIGH_PTR, dst_deepstack_ptr=MID_PTR, n_tokens=1369, hidden=3584, dtype="torch.bfloat16", has_deepstack=True, required_dst_info_num=1, ) frames = info.to_zmq() assert len(frames) == 13 # layout lock (v2: row_start + span appended) rt = EmbeddingTransferInfo.from_zmq(frames) assert rt == info assert rt.dtype == "torch.bfloat16" assert rt.has_deepstack is True def test_transfer_info_round_trip_no_deepstack(): info = EmbeddingTransferInfo( room=0, endpoint="::1", dst_port=65535, mooncake_session_id="x:0", dst_embedding_ptr=0, dst_deepstack_ptr=0, n_tokens=0, hidden=4096, dtype="torch.float32", has_deepstack=False, required_dst_info_num=2, ) rt = EmbeddingTransferInfo.from_zmq(info.to_zmq()) assert rt == info assert rt.has_deepstack is False def test_pointer_encoding_is_full_64bit(): # Pack/parse the boundary directly through a frame to ensure no truncation. info = EmbeddingArgsRegisterInfo( room=REGISTER_ROOM_SENTINEL, endpoint="h", dst_port=2, mooncake_session_id="s", dst_embedding_ptr=HIGH_PTR, dst_deepstack_ptr=0xDEAD_BEEF_CAFE_0000, ) rt = EmbeddingArgsRegisterInfo.from_zmq(info.to_zmq()) assert rt.dst_embedding_ptr == HIGH_PTR assert rt.dst_deepstack_ptr == 0xDEAD_BEEF_CAFE_0000 def test_transfer_info_integer_room_distinct_from_sentinel(): # The bootstrap thread routes on frame[0]: "None" => register, else int room. reg = EmbeddingArgsRegisterInfo(REGISTER_ROOM_SENTINEL, "h", 1, "s", 1, 0).to_zmq() xfer = EmbeddingTransferInfo( 7, "h", 1, "s", 1, 0, 10, 8, "torch.bfloat16", False, 1 ).to_zmq() assert reg[0] == b"None" assert xfer[0] == b"7" assert xfer[0] != b"None" # --------------------------------------------------------------------------- # # MooncakeEmbeddingSender (driven by a fake manager; no Mooncake engine) # --------------------------------------------------------------------------- # class _SenderFakeMgr: """Minimal stand-in for MooncakeEmbeddingManagerEncode: just the surface the sender touches.""" def __init__(self): self.request_status = {} self.failure_records = {} self.failure_lock = threading.Lock() self.bootstrap_time_out = 300 self.added = [] # (room, EmbeddingChunk) def update_status(self, room, status): self.request_status[room] = status def check_status(self, room): return self.request_status.get(room, TransferPoll.Bootstrapping) def record_failure(self, room, reason): self.failure_records[room] = reason def add_transfer_request(self, room, chunk): self.added.append((room, chunk)) def test_sender_failure_exception_raises_and_clears(): mgr = _SenderFakeMgr() s = MooncakeEmbeddingSender(mgr, "h:9", bootstrap_room=3) mgr.failure_records[3] = "boom on rank 1" with pytest.raises(EmbeddingTransferError) as exc_info: s.failure_exception() assert exc_info.value.bootstrap_room == 3 assert exc_info.value.failure_reason == "boom on rank 1" assert exc_info.value.remote_endpoint == "h:9" assert str(exc_info.value) == ( "EmbeddingTransferError(bootstrap_room=3, remote_endpoint=h:9): " "boom on rank 1" ) assert 3 not in mgr.request_status # cleared assert s.conclude_state == TransferPoll.Failed # --- shard_rows: the single source of shard geometry for both sides --- @pytest.mark.parametrize("span", [1, 2, 3, 7, 8, 128, 129]) @pytest.mark.parametrize("size", [1, 2, 4, 8]) def test_shard_rows_tiles_span_exactly(span, size): cursor = 0 for rank in range(size): start, count = shard_rows(span, rank, size) assert start == cursor # contiguous, in rank order assert count >= 0 cursor += count assert cursor == span # disjoint cover of [0, span) counts = [shard_rows(span, r, size)[1] for r in range(size)] assert max(counts) - min(counts) <= 1 # balanced # --- wire frame: round-trip + malformed input --- def _info(**overrides) -> EmbeddingTransferInfo: base = dict( room=7, endpoint="1.2.3.4", dst_port=5555, mooncake_session_id="s", dst_embedding_ptr=0x1000, dst_deepstack_ptr=0, n_tokens=10, hidden=64, dtype="torch.bfloat16", has_deepstack=False, required_dst_info_num=4, ) base.update(overrides) return EmbeddingTransferInfo(**base) # --- encode-side payload math + fanout validation --- def _chunk(n_tokens=10, hidden=64, deepstack_width=0) -> EmbeddingChunk: itemsize = 2 # bf16 return EmbeddingChunk( room=7, src_embedding_ptr=0x10_0000, n_tokens=n_tokens, hidden=hidden, dtype="torch.bfloat16", nbytes=n_tokens * hidden * itemsize, src_deepstack_ptr=0x20_0000 if deepstack_width else 0, deepstack_width=deepstack_width, deepstack_nbytes=( n_tokens * deepstack_width * itemsize if deepstack_width else 0 ), ) def test_shard_payload_identity_is_whole_chunk(): chunk = _chunk(n_tokens=10, deepstack_width=128) src, nbytes, deep_src, deep_nbytes = shard_payload(chunk, _info(n_tokens=10)) assert (src, nbytes) == (chunk.src_embedding_ptr, chunk.nbytes) assert (deep_src, deep_nbytes) == (chunk.src_deepstack_ptr, chunk.deepstack_nbytes) def test_shard_payload_offsets_rows(): chunk = _chunk(n_tokens=10, hidden=64, deepstack_width=128) row_bytes = chunk.nbytes // 10 deep_row_bytes = chunk.deepstack_nbytes // 10 src, nbytes, deep_src, deep_nbytes = shard_payload( chunk, _info(n_tokens=3, row_start=5) ) assert src == chunk.src_embedding_ptr + 5 * row_bytes assert nbytes == 3 * row_bytes assert deep_src == chunk.src_deepstack_ptr + 5 * deep_row_bytes assert deep_nbytes == 3 * deep_row_bytes def test_validate_rejects_gap_overlap_range_dtype(): chunk = _chunk(n_tokens=10) gap = [ _info(n_tokens=3, row_start=0, span=10), _info(n_tokens=3, row_start=5, span=10), ] overlap = [ _info(n_tokens=5, row_start=0, span=10), _info(n_tokens=5, row_start=3, span=10), ] out_of_range = [_info(n_tokens=6, row_start=5, span=10)] bad_dtype = [_info(n_tokens=10, dtype="torch.float32", span=10)] mixed = [_info(n_tokens=10, span=10), _info(n_tokens=5, row_start=5, span=10)] for infos in (gap, overlap, out_of_range, bad_dtype, mixed): assert validate_fanout_frames(infos, chunk) is not None def test_validate_rejects_span_mismatch_even_single_frame(): # The G2 token-count tripwire: a cross-side row-count divergence (image # processor / grid contract / cache-key bug) must fail LOUD even at # fanout == 1, where a lone frame has no tiling partner to expose it. chunk = _chunk(n_tokens=10) # span carries the receiver's full image span -> direct check. under = [_info(n_tokens=6, row_start=0, span=6)] assert validate_fanout_frames(under, chunk) is not None # matching span passes. assert validate_fanout_frames([_info(n_tokens=10, span=10)], chunk) is None def test_validate_rejects_deepstack_presence_mismatch(): # An encode chunk that lost its deepstack half (e.g. a cache hit cached # without it) must fail loud, not push Success while the receiver # publishes a never-written deepstack buffer; and vice versa. plain_chunk = _chunk(n_tokens=10) deep_chunk = _chunk(n_tokens=10, deepstack_width=128) wants_deep = [ _info(n_tokens=10, span=10, has_deepstack=True, dst_deepstack_ptr=0x2000) ] no_deep = [_info(n_tokens=10, span=10)] assert validate_fanout_frames(wants_deep, plain_chunk) is not None assert validate_fanout_frames(no_deep, deep_chunk) is not None assert validate_fanout_frames(wants_deep, deep_chunk) is None # --- receiver job: PER-ITEM shard placement + reassembly schedule --- class _ShardFakeEngine: def register(self, *a, **k): pass def deregister(self, *a, **k): pass class _ShardFakeMgr: def __init__(self): self.engine = _ShardFakeEngine() class _ShardReceiver: """Fake receiver: Bootstrapped until pre_alloc, Success after (mirrors test_encode_receiver's _FakeReceiver).""" created: list["_ShardReceiver"] = [] def __init__(self, manager, addr, room): self.addr = addr self.room = room self._pre_alloced = False self.pre_alloc_kwargs = None _ShardReceiver.created.append(self) def poll(self): return TransferPoll.Success if self._pre_alloced else TransferPoll.Bootstrapped def pre_alloc(self, **kwargs): self._pre_alloced = True self.pre_alloc_kwargs = kwargs def _shard_setup(monkeypatch): import tokenspeed.runtime.pd.epd.prefill_receiver as er _ShardReceiver.created.clear() # Pin the LEGACY per-request buffer path: the pointer-math assertions below # compare pre_alloc destinations against item.encoded, which on the pooled # path is a post-DONE CLONE at a different address. Pool+shard interplay is # covered separately by test_pool_shard_reassembles_into_published_clone. monkeypatch.setenv("TOKENSPEED_EPD_RECV_POOL_SLOTS", "0") er._POOLS.clear() def _shard_item(span, *, room=100, offsets=None): # ONE item == one EPD image of `span` concatenated-subgrid tokens, carrying # its per-item encode handshake. Multi-subgrid offsets still sum to `span` and # are received/sharded as one image over one room (the encode worker # concatenates the subgrids and row-splits the item's full embedding). item = MultimodalDataItem( modality=Modality.IMAGE, offsets=offsets if offsets is not None else [(0, span - 1)], ) item.encode_handshake = { "bootstrap_room": room, "bootstrap_host": "h", "bootstrap_port": 1, } return item def _start(items, shard_rank, shard_size, factory=_ShardReceiver, num_deepstack=0): return start_embedding_receive( items, manager=_ShardFakeMgr(), hidden=8, num_deepstack=num_deepstack, dtype=torch.float32, device="cpu", receiver_factory=factory, shard_rank=shard_rank, shard_size=shard_size, ) def test_packed_to_full_scatters_shard_rows_to_absolute_offsets(): # The publish scatter is the silent-corruption-critical mapping: packed shard # rows -> their absolute offsets in the item's full embedding. rank 1 of 2 on # a span-10 item: packed rows [0,5) -> full rows [5,10). Non-owned rows are # left for reassemble (asserted only on the owned rows here). import types job = _start([_shard_item(10)], shard_rank=1, shard_size=2) it = types.SimpleNamespace(n_tokens=10, spans=[10], row_starts=[5], row_counts=[5]) packed = torch.arange(5 * 8, dtype=torch.float32).reshape(5, 8) full = job._packed_to_full(it, packed, 8) assert full.shape == (10, 8) assert torch.equal(full[5:10], packed[0:5]) def _record_broadcasts(monkeypatch): calls = [] def fake_broadcast(tensor, src=None, group=None): calls.append((tensor, src)) monkeypatch.setattr(torch.distributed, "broadcast", fake_broadcast) return calls def test_reassemble_broadcasts_item_subranges(monkeypatch): # The reassembly schedule follows the item's shard tiling: each rank # broadcasts the contiguous row range it owns, together covering [0, span) # exactly once (one item == one image under per-item rooms). item = _shard_item(10) job = _start([item], shard_rank=0, shard_size=2) assert job.poll() == DONE calls = _record_broadcasts(monkeypatch) job.reassemble(nccl_group="g", group_ranks=(7, 9)) itemsize = item.encoded.element_size() base = item.encoded.data_ptr() rows = [ ((t.data_ptr() - base) // (8 * itemsize), t.shape[0], src) for t, src in calls ] # span 10, 2-way: rank0 rows [0,5) from global 7, rank1 rows [5,10) from 9. assert rows == [(0, 5, 7), (5, 5, 9)] # --- encode worker: concluded-sender sweep (the O1 hygiene rider) --- class _FakeReceiver: """Drives the poll state machine without any transport: Bootstrapped until pre_alloc, Success after. Records the pre_alloc kwargs for assertions.""" created: list["_FakeReceiver"] = [] def __init__(self, manager, addr, room): self.manager = manager self.addr = addr self.room = room self._pre_alloced = False self.pre_alloc_kwargs = None _FakeReceiver.created.append(self) def poll(self): return TransferPoll.Success if self._pre_alloced else TransferPoll.Bootstrapped def pre_alloc(self, **kwargs): self._pre_alloced = True self.pre_alloc_kwargs = kwargs class _RecvFakeEngine: """No-op Mooncake engine: receive_encoded_embeddings register/deregisters each receive buffer (a real RDMA NIC needs pre-registered targets); on CPU there is nothing to register.""" def register(self, *args, **kwargs): pass def deregister(self, *args, **kwargs): pass class _RecvFakeMgr: def __init__(self): self.engine = _RecvFakeEngine() def _recv_item(n_tokens: int) -> MultimodalDataItem: # _item_token_count sums (end - start + 1) over offsets; one subgrid here. return MultimodalDataItem(modality=Modality.IMAGE, offsets=[(0, n_tokens - 1)]) def _epd(item: MultimodalDataItem, *, room: int, host: str, port: int): """Attach an EPD encode->prefill handshake onto an item (one room per item).""" item.encode_handshake = { "bootstrap_room": room, "bootstrap_host": host, "bootstrap_port": port, } return item def _recv_setup(monkeypatch): import tokenspeed.runtime.pd.epd.prefill_receiver as er _FakeReceiver.created.clear() # Small pool defaults so each test's fresh fake engine doesn't allocate # the production 16x256MB region; pools are keyed by engine identity, so # clear them (and the lazy-dereg queue) between tests. monkeypatch.setenv("TOKENSPEED_EPD_RECV_POOL_SLOTS", "4") monkeypatch.setenv("TOKENSPEED_EPD_RECV_POOL_SLOT_MB", "1") er._POOLS.clear() er._pending_dereg.clear() def test_recv_pool_release_waits_for_clone_event(monkeypatch): import tokenspeed.runtime.pd.epd.prefill_receiver as er class _FakeEvent: def __init__(self): self.ready = False def query(self): return self.ready event = _FakeEvent() monkeypatch.setattr(er, "_record_current_stream_event", lambda _tensor: event) pool = er._RecvBufferPool(_RecvFakeEngine(), "cpu", slot_bytes=64, n_slots=1) slot = pool.lease(8) assert slot == 0 pool.release_after_copy(slot, torch.empty(1)) assert pool.lease(8) is None event.ready = True assert pool.lease(8) == 0 def test_receive_sizes_buffers_per_item_no_deepstack(): items = [ _epd(_recv_item(6), room=11, host="h0", port=7001), _epd(_recv_item(4), room=22, host="h1", port=7002), ] receive_encoded_embeddings( items, manager=_RecvFakeMgr(), hidden=2048, num_deepstack=0, dtype=torch.bfloat16, device="cpu", receiver_factory=_FakeReceiver, ) assert items[0].encoded.shape == (6, 2048) assert items[1].encoded.shape == (4, 2048) assert items[0].encoded.dtype == torch.bfloat16 assert items[0].encoded_deepstack is None assert items[1].encoded_deepstack is None # Each receiver pre_alloc'd a buffer matching its item's token count, with # the dtype string the encode side asserts against (str(torch.bfloat16)). r0 = next(r for r in _FakeReceiver.created if r.room == 11) assert r0.pre_alloc_kwargs["n_tokens"] == 6 assert r0.pre_alloc_kwargs["hidden"] == 2048 assert r0.pre_alloc_kwargs["dtype"] == "torch.bfloat16" assert r0.pre_alloc_kwargs["has_deepstack"] is False assert r0.pre_alloc_kwargs["dst_deepstack_ptr"] == 0 assert r0.addr == "h0:7001" def test_receive_allocates_deepstack_columns(): items = [_epd(_recv_item(5), room=99, host="h", port=8000)] receive_encoded_embeddings( items, manager=_RecvFakeMgr(), hidden=128, num_deepstack=3, dtype=torch.float32, device="cpu", receiver_factory=_FakeReceiver, ) assert items[0].encoded.shape == (5, 128) assert items[0].encoded_deepstack.shape == (5, 128 * 3) r = _FakeReceiver.created[0] assert r.pre_alloc_kwargs["has_deepstack"] is True assert r.pre_alloc_kwargs["dst_deepstack_ptr"] != 0 def test_receive_skips_already_encoded_item_on_recall(): # Chunked prefill re-runs receive_encoded_embeddings on the SAME item object # across forwards. After the first call sets item.encoded, a second call must # skip the item: construct NO new receiver (re-bootstrapping a room already at # Success would dead-wait phase-1 for Bootstrapped and time out) and leave the # encoded buffer untouched. item = _epd(_recv_item(5), room=77, host="h", port=9) common = dict( manager=_RecvFakeMgr(), hidden=16, num_deepstack=0, dtype=torch.float32, device="cpu", receiver_factory=_FakeReceiver, ) receive_encoded_embeddings([item], **common) assert item.encoded is not None assert len(_FakeReceiver.created) == 1 # first forward received it first_encoded = item.encoded # Second forward (chunked prefill): item.encoded is already set -> skipped. receive_encoded_embeddings([item], **common) assert len(_FakeReceiver.created) == 1 # no new receiver constructed assert item.encoded is first_encoded # buffer untouched