fix(maru_vllm): stop a fully cached prompt from killing the vLLM engine core

[youngrok-XCENA]

🤔 Background & Motivation (Why)

프롬프트의 토큰 수가 maru_kv_chunk_tokens 의 배수이고 그 프롬프트의 모든 청크가 풀에 있을 때, 이를 적재하려는 vLLM 인스턴스의 EngineCore 가 종료됩니다. 해당 인스턴스는 그 뒤로 아무 요청도 처리하지 못합니다. 프롬프트 토큰 수가 배수가 아니면 이 상황은 발생하지 않습니다.

Maru timing: lookup batch_exists 32 keys = 0.63 ms
Maru timing: get_num_new_matched (incl _chunk_keys) 8192 tok = 159.66 ms
ERROR EngineCore encountered a fatal error.
  File ".../vllm/v1/core/sched/scheduler.py", line 670, in schedule
    assert num_new_tokens > 0
AssertionError

프롬프트 토큰 수가 maru_kv_chunk_tokens 로 나누어떨어지면 청크 키가 프롬프트를 나머지 없이 덮습니다. 그러면 커넥터가 프롬프트 전체를 외부 캐시 히트로 보고하고, vLLM 스케줄러에는 이번 스텝에 계산할 토큰이 하나도 남지 않습니다. 부분 청크가 남는 프롬프트는 그 나머지가 언제나 일거리로 남기 때문에 영향을 받지 않습니다.

프롬프트	청크 256 기준	커넥터가 보고한 값	스케줄러에 남는 일	결과
8,192	32 청크 (나머지 0)	8,192	0	엔진 종료
32,768	128 청크 (나머지 0)	32,768	0	엔진 종료
8,210	32 청크 + 18 토큰	8,192	18	정상
3,858	15 청크 + 18 토큰	3,840	18	정상

이 조건 때문에 결함이 오래 눈에 띄지 않았습니다. 벤치마크는 입력 길이를 8k·32k 같은 둥근 값으로 잡는 일이 많아 매번 걸리는 반면, 길이가 제각각인 트래픽에서는 거의 나타나지 않습니다.

에러도 vLLM 내부 assert 로 나기 때문에 스택 트레이스 어디에도 maru 가 등장하지 않습니다. 원인을 모르면 vLLM 버그나 환경 문제로 읽힙니다.

단위 테스트로도 잡히지 않았습니다. TestDeferredLoading::test_matched_tokens_reported_async 가 64 토큰 프롬프트(8 청크 × 8 토큰, 나머지 0)에 대해 matched == 64, 즉 프롬프트 전체 보고를 정답으로 고정하고 있었습니다.

🏗️ Design Changes

스케줄러가 요청을 태우려면 그 요청에 엔진이 직접 계산할 몫이 남아 있어야 한다는 것이 vLLM 의 전제입니다. 커넥터가 프롬프트를 "풀에서 가져오는 몫"과 "엔진이 계산하는 몫"으로 가르는 주체이므로, 엔진 몫을 0 으로 만들지 않을 책임도 커넥터에 있습니다. 이 PR 은 그 책임을 커넥터가 지도록 바꿉니다.

flowchart LR
  subgraph before["Before"]
    direction TB
    BP["프롬프트<br/>(청크 크기의 배수)"]
    BPOOL["풀에서 가져오는 몫<br/>프롬프트 전체"]
    BENG["엔진이 계산하는 몫<br/>없음"]
    BDIE["스케줄러 전제 위반<br/>EngineCore 종료"]
    BP --> BPOOL
    BP --> BENG
    BENG --> BDIE
  end
  subgraph after["After"]
    direction TB
    AP["프롬프트<br/>(청크 크기의 배수)"]
    APOOL["풀에서 가져오는 몫<br/>마지막 한 블록을 뺀 전부"]
    AENG["엔진이 계산하는 몫<br/>한 블록"]
    AOK["정상 스케줄"]
    AP --> APOOL
    AP --> AENG
    AENG --> AOK
  end
  before ~~~ after

Loading

• 동작 변경: 외부 캐시 히트 보고량에 상한이 생깁니다. 프롬프트 전체를 히트로 보고하지 않습니다.
• 적용 범위: 프롬프트가 청크 크기의 배수일 때만 값이 달라집니다. 부분 청크가 있는 프롬프트는 이전과 같은 값을 보고합니다.
• 캐시 이득: 8,192 토큰 기준으로 마지막 한 블록(16 토큰)만 재계산하므로 이득의 99.8% 가 유지됩니다.
• 적재 경로는 바꾸지 않았습니다. 워커는 여전히 캐시된 청크를 전부 읽어 오고, 겹치는 마지막 블록은 엔진이 재계산하며 덮어씁니다. 같은 값이라 무해합니다.
• 저장 경로, 키 구성, 청크 크기, 와이어 프로토콜은 그대로입니다.

📝 Implementation Details

근본 원인은 MaruSchedulerConnector.get_num_new_matched_tokens 가 캐시된 청크 수를 토큰 수로 환산한 값을 그대로 돌려준 데 있습니다.

matched_tokens = num_matched_chunks * self._kv_chunk_tokens
matched_tokens = _align_down(matched_tokens, self._block_size)
new_matched = matched_tokens - num_computed_tokens

num_matched_chunks * kv_chunk_tokens 가 len(token_ids) 와 같아질 수 있고, 그때 _align_down 도 값을 줄이지 않습니다(프롬프트 길이가 청크 크기의 배수면 블록 크기의 배수이기도 하므로). vLLM 쪽에서는 이렇게 이어집니다.

# vllm/v1/core/sched/scheduler.py
num_new_tokens = request.num_tokens - num_computed_tokens
...
assert num_new_tokens > 0

수정은 정렬 앞에 상한을 한 줄 넣는 것입니다.

matched_tokens = min(matched_tokens, len(token_ids) - 1)

_align_down 이 뒤따르므로 실제로는 블록 하나가 통째로 남습니다. 8,192 토큰 · 블록 16 이면 min 이 8,191 로 깎고 정렬이 8,176 으로 내려, 엔진이 마지막 16 토큰을 계산합니다.

num_matched_chunks 는 줄이지 않았습니다. 보고량을 줄이면서 적재량까지 줄이면 경계 계산이 두 군데로 갈라지는데, 겹치는 마지막 블록은 엔진이 어차피 같은 값으로 덮어쓰므로 얻는 것이 없습니다.

리뷰에서 봐 주셨으면 하는 지점은 상한의 위치입니다. _align_down 뒤에 두면 len(token_ids) - 1 이 블록 정렬을 깨서 보고량이 블록 경계에 맞지 않게 됩니다. 앞에 두어야 정렬이 마지막으로 적용됩니다. 회귀 테스트에 정렬 유지 조항을 따로 두었습니다.

✅ Tests

• Unit tests
• Integration tests
• Manual tests
• No tests needed (reason: )

Unit. TestFullPromptHitLeavesWorkForTheEngine 을 추가했습니다. 청크 크기의 배수인 프롬프트(64 / 128 / 256 토큰)에 대해 보고량이 프롬프트 길이보다 작은지, 그리고 블록 정렬이 유지되는지를 확인하고, 부분 청크가 있는 프롬프트의 동작이 그대로인지도 함께 봅니다. 수정 전에는 3 건이 실패하고 수정 후에는 통과합니다. 비통합 스위트 전체 910 통과 4 스킵입니다.

기존 test_matched_tokens_reported_async 는 프롬프트를 64 토큰에서 70 토큰(8 청크 + 나머지)으로 바꿨습니다. 원래 의도인 "캐시된 청크는 전부 보고한다"는 그대로 검증하면서, 새로 생긴 상한 규칙과 겹치지 않게 했습니다.

packed 형식 왕복 커버리지도 함께 넣었습니다. 이 결함과는 별개로 비어 있던 자리입니다. 기존 왕복 테스트(TestStoreLoadRoundtripModernLayout)는 layerwise 형식을 레이어 한 장으로만 돌고, 기본값인 packed 형식은 저장한 슬랩을 다시 읽어 값을 대조하는 테스트가 없었습니다.

• TestPackedStoreLoadRoundtrip — 저장된 슬랩의 바이트 배열이 선언한 [2, num_layers, chunk_tokens, hidden] 과 맞는지를, 커넥터 자신의 판독기가 아니라 테스트가 따로 만든 판독기로 대조합니다. 왕복만 보면 저장·적재 양쪽이 같은 잘못된 규약을 공유할 때 통과해 버리는데, 생산자와 소비자가 서로 다른 프로세스인 상황이 정확히 그 경우입니다. 이어서 다른 인스턴스가 읽어 복원한 값이 원본과 같은지도 확인합니다. CPU 에서 돌아가므로 per-layer fallback 경로를 덮습니다.
• TestPackedKernelRoundtripGPU — CUDA 와 lmcache.c_ops 가 있을 때만 수행합니다. 페이지 고정 호스트 메모리를 풀 대역으로 세워(CXL 장치 불필요) fused 커널 경로를 돌고, 기본 저장 경로와 maru_async_store 의 GPU staging 저장 경로가 바이트 단위로 같은 객체를 만드는지 대조합니다.

두 묶음 모두 변이 주입으로 판별력을 확인했습니다. 적재 쪽 레이어 인덱스를 어긋내면 왕복만 실패하고 바이트 배열 테스트는 통과하며, 저장 쪽을 어긋내면 둘 다 실패합니다. staging 경로에만 레이어 축 이동을 주입하면 해당 경로의 테스트만 실패합니다.

Manual (E2E). vLLM 0.22.1, Qwen2.5-7B (28 layers, bf16), TP1 × 2 (RTX PRO 6000 Blackwell), GPU 당 KV 예산 16 GiB, 엔진 접두사 캐시 끔, maru_kv_chunk_tokens=256, maru 풀 4 GiB. 프롬프트는 토큰 단위로 정확히 8,192 토큰 4 종, 출력 128 토큰, 동시성 1, 워밍업을 분리한 뒤 16 건 측정.

구성	수정 전	수정 후
인스턴스 2 대가 풀 공유 (한쪽이 저장, 다른 쪽이 적재)	적재 측 EngineCore 종료	16 건 전부 저장 측과 같은 128 토큰, TTFT 중앙값 51.8 ms
프리필·디코드 분리 (프리필은 1 토큰만 생성)	같은 원인으로 불가	16 건 전부 기준 답과 일치, TTFT 중앙값 52.7 ms

수정 후 적재 측 TTFT 중앙값은 51.8 ms 로, 같은 프롬프트를 처음부터 계산할 때의 465 ms 대비 약 9 배 빠릅니다. 두 구성 모두 적재 측에서 매 요청 32 개 키를 조회하고 packed-load kernel 28L x 32c 가 기록됩니다.

🔗 Related Issues (optional)

📦 Release Note (for auto-generation / write in English)

NEW

CHANGED

• maru_vllm: added store/load round-trip tests for the packed (default) chunk format, including a GPU check that the fused kernel and the maru_async_store staging path write identical objects.

FIXED

• maru_vllm: a prompt whose token count is an exact multiple of maru_kv_chunk_tokens no longer kills the vLLM engine core on a cache hit. The connector now reports at most prompt length - 1 tokens as an external hit, so the scheduler always has a block left to compute.

IMPORTANT NOTES

[github-actions]

Coverage Report

File	Stmts	Miss	Cover	Missing
__init__.py	6	0	100%
__main__.py	3	3	0%	5, 7–8
allocation_manager.py	129	11	91%	30–31, 44–45, 52, 237–238, 268–271
client.py	187	8	96%	88, 130, 147–148, 311–313, 319
config.py	46	1	98%	72
constants.py	9	0	100%
device_scanner.py	115	5	96%	13–15, 63, 96
handler.py	557	86	85%	153, 168, 179–186, 194–196, 211, 217–219, 227, 238–243, 248, 273, 300–301, 305, 345, 349–350, 355, 377, 384–385, 389–395, 398, 402–404, 420, 422–423, 432, 486, 507, 517–518, 624–626, 633, 755–758, 761, 772–775, 781, 867, 971, 1134–1138, 1144, 1155–1159, 1165, 1244, 1249, 1291
ipc.py	275	2	99%	365, 441
kv_manager.py	108	0	100%
logging_setup.py	19	0	100%
plugin.py	59	3	95%	65–67
protocol.py	233	0	100%
resource_manager_installer.py	103	13	87%	80–86, 167, 169–172, 187
rpc_async_client.py	190	0	100%
rpc_async_server.py	111	0	100%
rpc_client.py	71	0	100%
rpc_client_base.py	106	10	91%	185, 221–222, 233–234, 306–307, 311–312, 375
rpc_handler_mixin.py	106	19	82%	154–156, 159–161, 219–222, 227, 231–235, 256–257, 263
rpc_server.py	64	0	100%
serializer.py	81	0	100%
server.py	170	13	92%	64–65, 73, 168, 172, 243, 247, 265–266, 357–359, 444
stats_manager.py	95	0	100%
types.py	60	1	98%	145
uds_helpers.py	13	0	100%
memory
__init__.py	5	0	100%
allocator.py	55	0	100%
mapper.py	164	17	90%	65–74, 77–80, 293, 342, 376
owned_region_manager.py	101	1	99%	212
types.py	63	0	100%
TOTAL	3325	193	94%

Tests	Skipped	Failures	Errors	Time
724	4 💤	0 ❌	0 🔥	7.903s ⏱️

[kihwan-XCENA]

작성해주신 코드 확인했습니다.

전체 토큰이 히트되어도 -1은 해줘야 하는군요..

LGTM!