feat(maru_vllm): run the load/compute overlap on the layerwise storage format - #78
feat(maru_vllm): run the load/compute overlap on the layerwise storage format#78kihwan-XCENA wants to merge 5 commits into
Conversation
…e format maru_overlap_load_with_compute previously required the packed layout and was silently disabled under maru_use_layerwise=true. Relax the two _layerwise_overlap gates and the worker submit gate so layerwise-format deferred loads also go through the off-thread loader, keeping the safety invariant: every (chunk, layer) key resolves before the unpark gate can fire, and any miss fails the request while it is still parked so vLLM recomputes. The retained payload is keyed by true layer index (the loader thread and the resumed forward walk layers in different orders), and a new _copy_layerwise_layer_to_device gathers one layer's per-chunk objects with run-coalesced contiguous H2D copies — no pitched copy needed, since layerwise objects are already single-layer. The unpark gate, per-layer events, wait_for_layer_load, and preemption drain are reused unchanged. Overlap-off layerwise requests keep the existing sync path. Also fix the mislabeled key-count figures in the docs (59 vs 1,888 was a ~15k prompt; a 64k prompt resolves ~250 vs ~8,000).
Review follow-ups on the layerwise-overlap change, all within its scope: - Reject a batch_retrieve response shorter than the key list, in both the off-thread job and the sync path. A truncated response passed the None-only miss scan; under the layerwise layout it would shift every later layer's objects by one position, and the staging gather would inject uninitialized device memory for the missing tail — silent wrong KV, now a logged recompute while the request is still parked. The gather helper also refuses fewer objects than chunks outright. - Log the off-thread job's mismarked-request guard at error level; it previously degraded to recompute with no diagnostic. - Extract _chunk_layer_key/_load_keys so the store, the sync load, and the off-thread load share one owner of the _L<idx> key format. - Make the retained-payload annotations tell the truth (list-or-dict), use element_size() like the rest of the file, neutralize the "packed-layerwise" log strings, and fix the two remaining stale key-count figures in _load_packed's comments.
The docs table and the maru_use_layerwise knob description already name the default format chunkwise; the prose this branch added said packed for the same thing. Align the comments, docstrings, and docs on chunkwise. Identifiers keep their existing names.
The key-count comparison was written as a worked example for one model and one chunk size, but the chunk size is configurable and the layer count varies by model, so the numbers only held for the defaults. State the relationship instead: layerwise costs one key per (chunk, layer) rather than one per chunk, so both the key count and the retrieve metadata RPC volume scale with the model's layer count.
…handles Three helpers on the deferred-load path became format-agnostic when the overlap learned the layerwise format, but their names still said packed: drop it from _deferred_load_job, _try_submit_deferred_load, and _schedule_deferred_layerwise_loads. Rename _copy_packed_layer_to_device to _copy_chunkwise_layer_to_device so the pair selected by maru_use_layerwise reads as chunkwise vs layerwise, matching the vocabulary the docs already use. The remaining chunkwise-only helpers keep their packed names for a separate cleanup.
youngrok-XCENA
left a comment
youngrok-XCENA
left a comment
There was a problem hiding this comment.
LGTM. 두 손잡이를 독립시키는 방향이 옳고, 구현도 그 독립성을 실제로 지킵니다. 특히 세 가지가 좋았습니다.
첫째, retrieve 응답이 키 개수보다 짧을 때의 무증상 경로를 찾아낸 점입니다. 기존 검사는 빈 슬롯만 훑었기 때문에 짧은 응답은 통과했고, 결과를 위치로 슬라이싱하는 구조상 이후 모든 층이 앞 층의 객체를 읽으면서 staging 버퍼 꼬리는 초기화되지 않은 채 모델에 들어갔습니다. 세 지점 모두에서 개수를 먼저 비교하도록 한 수정은 이 PR의 본 기능과 별개로 그 자체로 가치가 큽니다.
둘째, loader 스레드와 재개된 forward의 층 순회 순서가 다를 수 있다는 점을 위치 리스트가 아니라 실제 층 인덱스를 키로 하는 맵으로 해결한 점입니다. 순서 가정이 코드에 남아 있지 않습니다.
셋째, submit gate가 절대 보내지 않는 조합이 off-thread job에 도달했을 때 조용히 실패하지 않고 배선 오류로 로그를 남기는 점입니다. 이런 실패는 캐시 적중률 하락으로만 보여서 추적이 어려운데, 진단 문구가 함께 있습니다.
새 테스트 15건도 이 경로들을 실제로 덮고 있어 확인이 수월했습니다.
병합을 막지 않는 후속 항목 세 가지를 남깁니다.
-
_active_deferred_req_ids(connector.py:703, 884, 1053)는 쓰기만 하고 읽는 곳이 없습니다. 이 PR이 그 주석을 "layerwise overlap이 실제로 singleton 워크로드일 때만 쓰이도록"으로 유지했는데, connector.py:908이layerwise_load = self._layerwise_overlap이라 실제로는 모든 deferred load가 overlap을 탑니다.MaruReqMeta.layerwise_load주석은 정확히 고쳐졌으니 이쪽도 맞추거나, 읽는 곳이 없는 필드 자체를 정리하면 좋겠습니다. -
_chunk_layer_key를_L<idx>접미사의 단일 소유자로 도입한 취지에 맞추면,_DONE접미사도 connector.py:741과 connector.py:3611 두 곳의 문자열 리터럴로 남아 있으니 같은 방식으로 묶을 수 있습니다. -
_copy_layerwise_layer_to_device는 CXL 매핑에서 만든torch.frombuffer텐서를copy_(non_blocking=True)의 소스로 씁니다. 그 영역이 CUDA에 등록되어 있지 않으면 non-blocking 지정과 무관하게 host 동기 복사가 되어 loader 스레드가 층마다 직렬화됩니다. 측정에서 이득이 나왔으니 실제로는 등록되어 있을 것으로 보입니다만, chunkwise 경로가 "registered CXL"을 전제로 명시해 둔 것처럼 이쪽에도 적어두면 나중에 성능 회귀를 오진하지 않습니다.
병합 순서 관련 한 가지. #71과 docs/source/integration/vllm.md가 충돌합니다. connector.py는 깨끗하게 합쳐집니다. 다만 #71이 추가한 "overlap의 전제는 chunkwise 저장(maru_use_layerwise를 false로 유지)" 문단이 이 PR 병합 이후 사실과 어긋나므로, 나중에 병합되는 쪽에서 그 문단까지 정리해야 합니다.
Jenkins pr-head가 아직 대기 중이니 머지는 그 뒤에 부탁드립니다.
2 participants
🤔 Background & Motivation (Why)
maru_overlap_load_with_computehides most of a cache hit behind the request's own prefill: the per-layer copies are queued while the request waits, and the request starts as soon as layer 0 lands, so the rest of the layers arrive while earlier ones are already computing.This only worked with the chunkwise storage format. Under
maru_use_layerwise=truethe connector turned the overlap off with a warning, and a layerwise cache hit did its key lookup on the engine thread, inside the forward pass.The two knobs are meant to be independent —
maru_use_layerwisepicks how KV is grouped into CXL objects,maru_overlap_load_with_computepicks how a load is pipelined. This PR makes them independent:🏗️ Design Changes
Behavioral: either storage format can use the overlap. A layerwise cache hit now does its key lookup on the loader thread instead of the engine thread, so it no longer blocks the forward pass.
Behavioral: a lookup that returns fewer objects than keys is now rejected. The batch lookup answers one entry per key, in order, with a miss marked empty — and the old check only scanned for those empty slots. A reply that was simply shorter than the key list passed it, and because the results are then sliced by position, every layer after the gap read the previous layer's objects, while the tail of the GPU staging buffer was never written and reached the model as uninitialized KV. Both were silent. Every load path now compares the two counts first, logs a mismatch, and recomputes the request.
Dependency: layerwise no longer touches LMCache. A chunkwise object holds every layer of a chunk, so reading one layer means slicing it out of each object — that is what the borrowed LMCache kernel is for. A layerwise object is already one layer, stored end to end, so ordinary copies plus PyTorch's scatter are enough. Layerwise deployments therefore do not depend on the installed LMCache version at all.
Structural: retrieved objects are keyed by layer index. The loader thread and the resumed forward walk the layers in orders that need not match, so what the loader keeps for the forward is a map from layer index to that layer's objects, not a positional list.
📝 Implementation Details
✅ Tests
Unit tests — 15 new: the gates, key building, "a miss, or a reply short of the key count, fails while the request waits", CPU fallback round trip, per-layer CUDA events, an overlap-off regression, and a CUDA end-to-end (background load → release → per-layer waits, bit-exact KV). Whole suite: 913 passed on a 2-GPU node, ruff clean.
Manual tests — real node, one config run with the overlap off and on (naru
cfg/p2p/long_doc_singlenode_maru_direct_layerwise_overlap.yaml):kv_connector_extra_config— the two arms differ only in the last line:Overlap off → on, on the querying instance. These are cache-hit requests — the overlap only acts on a load, so a miss looks the same either way. vLLM's own prefix cache is off, so the same prompts computed cold took 2,617 ms (1 in flight) and 8,657 ms (8 in flight) to first token.
examples/vllm/p2p_sharing— a request answered from a layerwise-overlap cache hit produced byte-identical text to the same prompt computed from scratch, and the logs confirm it went through the new path rather than falling back.📦 Release Note (for auto-generation / write in English)
NEW
maru_overlap_load_with_computenow works with the layerwise storage format (maru_use_layerwise=true); it only needsmaru_async_load=true.CHANGED
FIXED
batch_retrievereply shorter than the key list is rejected and recomputed instead of loading uninitialized or layer-shifted KV.IMPORTANT NOTES