From bb041f6b1ae3cb42b063b0cf43134f03140a7329 Mon Sep 17 00:00:00 2001 From: rongchenghao Date: Sun, 9 Aug 2026 14:41:25 +0800 Subject: [PATCH] [perf] Copy via memoryview in pack_into instead of torch.frombuffer pack_into built a uint8 tensor per item with torch.frombuffer and dispatched a copy_ kernel for it. At TQ's object sizes the dispatch dominates: one key is stored per (field, sample), so items are a few hundred bytes each and the tensor construction plus kernel launch costs far more than the memcpy. Copy through memoryview slice assignment instead. Casting both sides to "B" keeps this correct whether the caller passes bytes, a memoryview or a numpy array. Also compute the size requirement from the item views already materialised for the copy, so calc_packed_size no longer walks every item a second time. Measured (Ascend NPU container, torch 2.9.0+cpu): 1024 objects x 512B: 35.0ms -> 5.5ms (6.4x) 7168 objects x 512B: 248.6ms -> 40.7ms (6.1x) 2048 objects x 16KB: 82.5ms -> 22.5ms (3.7x) In a real verl GRPO run, pack_into for a 512-sample 2-field put went 51.3ms -> 6.0ms, taking put_data from 183.9ms to 139.7ms single-node and 364.9ms to 290.4ms dual-node. Output is byte-for-byte identical to the previous implementation; both serial_utils test modules pass, and the packed buffers still round-trip through unpack_from/decode unchanged. Note both callers benefit: YuanrongStorageClient and MooncakeStore both go through batch_encode_into. The measurements above are from the yuanrong path; Mooncake was not measured. --- transfer_queue/utils/serial_utils.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/transfer_queue/utils/serial_utils.py b/transfer_queue/utils/serial_utils.py index e8ab267a..2ba8e7d1 100644 --- a/transfer_queue/utils/serial_utils.py +++ b/transfer_queue/utils/serial_utils.py @@ -438,23 +438,23 @@ def calc_packed_size(items: Sequence[bytestr]) -> int: def pack_into(target_buffer: bytestr, items: Sequence[bytestr]) -> None: """Concatenate ``items`` into ``target_buffer``, which must be at least ``calc_packed_size(items)`` bytes.""" - target_mv = memoryview(target_buffer) - required = calc_packed_size(items) + target_mv = memoryview(target_buffer).cast("B") + # Materialise the views once: they are needed for both the size check and the + # copy, and calc_packed_size would otherwise walk every item a second time. + item_mvs = [memoryview(item).cast("B") for item in items] + count = len(item_mvs) + required = _PACK_HEADER_SIZE + count * _PACK_ENTRY_SIZE + sum(mv.nbytes for mv in item_mvs) if target_mv.nbytes < required: raise ValueError(f"pack_into: target buffer has {target_mv.nbytes} bytes, requires {required}") - struct.pack_into(_PACK_HEADER_FMT, target_mv, 0, len(items)) + struct.pack_into(_PACK_HEADER_FMT, target_mv, 0, count) entry_offset = _PACK_HEADER_SIZE - payload_offset = _PACK_HEADER_SIZE + len(items) * _PACK_ENTRY_SIZE + payload_offset = _PACK_HEADER_SIZE + count * _PACK_ENTRY_SIZE - target_tensor = torch.frombuffer(target_mv, dtype=torch.uint8) - - for item in items: - item_mv = memoryview(item) + for item_mv in item_mvs: nbytes = item_mv.nbytes struct.pack_into(_PACK_ENTRY_FMT, target_mv, entry_offset, payload_offset, nbytes) - src_tensor = torch.frombuffer(item_mv, dtype=torch.uint8) - target_tensor[payload_offset : payload_offset + nbytes].copy_(src_tensor) + target_mv[payload_offset : payload_offset + nbytes] = item_mv entry_offset += _PACK_ENTRY_SIZE payload_offset += nbytes