fix: per-layer store copy, test plugin isolation, async-load block budget

[youngrok-XCENA]

🤔 Background & Motivation (Why)

이 PR 은 서로 독립된 결함 세 가지를 고칩니다. 결함마다 커밋이 있어 커밋 단위로 리뷰할 수 있습니다.

변경(커밋)	결함	영향
1. 층별 저장 복사(5e89c8a)	CPU·mixed 모드의 층별 저장 복사가 비연속 목적지에 한 번에 써서, torch CPU 스레드가 많을 때 느려짐	vLLM 서빙에서 저장하는 요청의 TTFT 중앙값 73.6 ms, 최대 179.3 ms(고친 뒤 중앙값 53.0 ms, 최대 58.2 ms)
2. 시험의 플러그인 격리(9d9486c)	단위 시험이 설치된 장치 플러그인을 불러 실제 DAX 장치에 힌트를 보냄	시험이 다른 작업의 장치를 건드림
3. 비동기 적재 대기의 블록 예산(5a5a66d, 4c7e6a6, b6f20fb)	maru_async_load 에서 적재를 기다리는 요청들이 GPU KV 블록을 나눠 쥐어, 가장 먼저 온 요청이 나머지 블록을 받지 못함	vLLM 엔진이 아무 요청도 처리하지 않고 멈춤. 클라이언트가 요청을 끊어야 풀림

1. 층별 저장 복사가 느림

CPU·mixed 모드(#88)에서 KV 를 저장하는 층별 복사가 비연속 목적지에 한 번에 씁니다. 이 복사는 torch CPU 스레드 수가 많으면 크게 느려집니다.

• CPU·mixed 모드의 vLLM 커넥터는 스텝 안에서 GPU KV 를 핸들러가 내준 호스트 버퍼의 슬랩으로 동기 복사합니다. 슬랩은 한 청크의 모든 층 K·V 를 이어 붙인 버퍼입니다.
• 슬랩은 [K·V(2), 층, 토큰, hidden] 순서로 배치됩니다. 그래서 한 층의 K 평면과 V 평면 사이에는 다른 층들이 끼어 있습니다.
• 현재 코드는 층마다 그 층의 K·V 를 슬랩의 그 층 자리에 복사 한 번으로 넣습니다. 이 목적지는 연속 메모리가 아닙니다.
• vLLM 없이 커넥터만 떼어 쟀습니다. Qwen2.5-0.5B 형상(24층)에서 3,304 토큰 프롬프트의 청크 12개(3,072 토큰)를 저장하는 시간은 torch CPU 스레드 1·8개에서 19 ms 이고, 128개(측정 호스트의 코어 수이자 torch 기본값)에서 1.76–1.82 s 입니다. 원인은 비연속 목적지 복사를 torch CPU 스레드 풀이 나눠 처리하는 비용으로 추정하며, 프로파일로 확인하지는 않았습니다.
• 저장은 동기라서 저장 시간은 그 스텝에 묶인 모든 요청의 스텝 시간에 더해집니다.
• vLLM 서빙에서는 저장하는 요청의 TTFT(요청을 보낸 때부터 첫 토큰을 받은 때까지)가 중앙값 73.6 ms, 최대 179.3 ms 입니다(20회 중 6회가 100 ms 초과). 이 TTFT 는 위 커넥터 단독 저장 시간과 다른 양입니다. vLLM 서빙이 쓰는 torch 스레드 수는 확인하지 않았습니다.
• 이 경로는 CPU 와 mixed 모드가 항상 쓰고, CXL 모드도 maru_kv_ops 커널을 쓰지 못할 때 씁니다.

2. 단위 시험이 설치된 장치 플러그인을 불러옴

시험 환경에 장치 벤더 플러그인이 설치되어 있으면, 단위 시험이 만든 핸들러가 그 플러그인을 불러 실제 DAX 장치에 힌트를 보냅니다.

• MaruHandler 는 생성될 때 maru.handler_plugins 엔트리포인트에 등록된 플러그인을 모두 불러옵니다.
• 단위 시험의 가짜 공유 메모리 클라이언트는 영역의 장치 경로를 /dev/dax0.0 으로 알려 줍니다. 플러그인은 이 경로를 실제 장치로 보고 힌트를 보내며, 그 장치는 같은 호스트의 다른 작업이 쓰고 있을 수 있습니다.
• 개발자 셸의 MARU_PLUGINS 값도 시험에 그대로 들어갑니다. 예를 들어 플러그인을 끄려고 MARU_PLUGINS=none 을 내보내면 플러그인 로더 시험 4개가 실패합니다.

3. 비동기 적재 대기 요청들이 GPU KV 를 모두 쥐고 멈춤

maru_async_load 를 켜면, 적재가 오래 걸릴 때 vLLM 엔진이 아무 요청도 처리하지 않고 멈출 수 있습니다. 클라이언트가 요청을 끊어 블록이 풀려야 다시 움직입니다.

• 비동기 적재를 쓰면 vLLM 은 Maru 에 접두사가 있는 요청을 적재 대기 상태(WAITING_FOR_REMOTE_KVS)에 둡니다. 이때 블록은 적재할 접두사 몫만 잡습니다.
• 받을 때 vLLM 은 그 요청의 프롬프트 전체가 그 순간의 빈 블록에 들어가는지 확인합니다(scheduler_reserve_full_isl, 기본 켬). 요청마다 그 순간의 빈 블록과 비교하므로, 적재 대기 요청이 여럿 쌓이면 그 요청들이 다 함께 필요로 하는 블록은 GPU KV 전체를 넘을 수 있습니다.
• 적재가 끝나면 가장 먼저 온 요청부터 나머지 블록을 받습니다. 그 몫이 빈 블록보다 크면 스케줄러는 그 스텝에서 대기 요청을 더 받지 않습니다. 실행 중인 요청이 없으면 블록을 내놓을 요청도 없어, 다음 스텝도 같습니다.
• 적재가 오래 걸릴수록 대기 요청이 더 많이 쌓이고, 프롬프트가 길수록 그 요청들이 다 함께 필요로 하는 블록이 커집니다.
• 실제 vLLM 0.22.1 스케줄러로 재현했습니다. GPU 블록 64개에 프롬프트 전체가 50·25·15블록인 요청을 차례로 적재 대기에 넣으면, 셋 다 받아들여져 62블록을 쥐고, 빈 블록은 1개만 남습니다(64개 중 1개는 vLLM 이 null 블록으로 따로 둡니다). 적재를 끝낸 뒤 첫 요청은 나머지 26블록을 받지 못하고, 다음 스텝에 아무 요청도 실행되지 않습니다.
• 실제 서빙에서도 재현했습니다. GPU 블록 64개인 엔진에 800·400 토큰 요청을 함께 보내면 main 은 3회 모두 멈췄습니다(아래 Tests).

🏗️ Design Changes

1. 층별 저장 복사

한 층의 K 와 V 를 각각 연속 영역으로 두 번 복사합니다. 슬랩의 배치와 바이트는 그대로이고, 복사하는 단위만 바뀝니다.

flowchart LR
    subgraph before["Before: 층마다 복사 한 번"]
        g1["GPU 의 한 층 K·V"] -->|"복사 한 번, 비연속 목적지"| s1["슬랩의 그 층 자리<br/>K 평면과 V 평면이 떨어져 있음"]
    end
    subgraph after["After: 층마다 복사 두 번"]
        g2["GPU 의 한 층 K"] -->|"연속 복사"| k2["슬랩 K 평면의 그 층"]
        g3["GPU 의 한 층 V"] -->|"연속 복사"| v2["슬랩 V 평면의 그 층"]
    end

Loading

2. 시험의 플러그인 격리

모든 시험에서 플러그인 엔트리포인트 탐색이 빈 목록을 돌려주고 플러그인 허용 목록 환경 변수가 비워집니다. 플러그인 로더 시험은 원래대로 자기 엔트리포인트를 직접 넘깁니다.

3. 비동기 적재 대기의 블록 예산

스케줄러 쪽 커넥터가 적재 대기 요청마다 프롬프트 전체에 필요한 블록 수를 기록합니다. 새 요청을 적재 대기로 받으면 그 합이 요청이 쥘 수 있는 GPU KV 블록 수를 넘게 되는 경우, vLLM 에 "다음 스텝에 다시 물어 달라"고 답합니다. 이렇게 하면 실행 중인 요청이 없어도 가장 먼저 온 대기 요청이 들어갈 자리가 늘 남습니다. 먼저 기다리게 한 요청의 몫은 뒤에 온 요청보다 먼저 잡아 두어, 긴 프롬프트가 계속 밀리지 않게 합니다. vLLM 은 기다리게 한 요청을 매 스텝 다시 묻는데, 예산에 자리가 날 때까지는 Maru 에 접두사를 다시 조회하지 않고 예산만 보고 답합니다.

flowchart TB
    subgraph before["Before"]
        b1["접두사가 있는 요청"] --> b2["그 순간 빈 블록에 프롬프트 전체가 들어가면 받음<br/>(vLLM 의 확인)"]
        b2 --> b3["적재 대기 요청이 쌓여<br/>다 함께 필요한 블록이 GPU KV 를 넘을 수 있음"]
        b3 --> b4["첫 요청이 나머지 블록을 못 받고 멈춤"]
    end
    subgraph after["After"]
        a1["접두사가 있는 요청"] --> a2{"적재 대기 요청들과<br/>먼저 기다린 요청들의 몫 + 이 요청<br/>≤ GPU KV 블록 수?"}
        a2 -->|"예"| a3["적재 대기로 받음"]
        a2 -->|"아니오"| a4["다음 스텝에 다시 물음<br/>(먼저 기다린 순서 유지,<br/>자리가 날 때까지 조회 없음)"]
        a3 --> a5["적재 뒤 첫 요청이 늘 들어감"]
    end

Loading

📝 Implementation Details

1. 층별 저장 복사

• maru_vllm/connector.py MaruWorkerConnector._save_kv_layer_packed 의 커널 없는 층별 복사:
  • slab[:, layer_idx].copy_(kv_contig) 를 평면마다 slab[plane, layer_idx].copy_(kv_contig[plane]) 로 바꿨습니다.
  • docstring 과 주석도 함께 고쳤습니다.
• 기동 조건, 설정, 저장 형식은 바뀌지 않습니다.

2. 시험의 플러그인 격리

• tests/conftest.py 에 autouse 픽스처 _no_installed_handler_plugins 를 추가했습니다. maru_handler.plugin._discover_entry_points 를 빈 목록으로 바꾸고 MARU_PLUGINS 를 지웁니다.
• tests/unit/test_plugin_loader.py 에 시험 안의 핸들러가 플러그인을 하나도 불러오지 않는지 보는 시험을 추가했습니다.

3. 비동기 적재 대기의 블록 예산

• maru_vllm/connector.py:
  • MaruKVConnector 가 스케줄러 쪽 MaruSchedulerConnector 에 요청이 쥘 수 있는 GPU 블록 수를 넘깁니다(_usable_gpu_blocks(): kv_cache_config.num_blocks, 없으면 cache_config.num_gpu_blocks 에서 vLLM 의 null 블록 1개를 뺀 값, 4c7e6a6).
  • get_num_new_matched_tokens() 의 기존 일치 계산을 _match() 로 떼어 내고, 그 결과가 비동기 적재일 때 _admit_parked() 로 예산을 확인하고, 넘치면 (None, False) 를 돌려줍니다. vLLM 은 이 답을 받으면 그 요청을 다음 스텝에 다시 묻습니다.
  • 이미 기다리는 요청을 다시 물으면 _admit_parked() 를 먼저 보고, 자리가 없으면 _match() 를 건너뛰고 (None, False) 를 돌려줍니다(b6f20fb). 비동기 적재는 CXL 모드에서만 켤 수 있습니다. 이 모드의 _match() 는 프롬프트의 청크 가운데 아직 확인하지 않은 키가 있으면 스케줄러 쪽 핸들러로 MaruServer 에 존재 확인을 보내고, 적중하면 "Maru KV hit" INFO 로그를 남깁니다. 그래서 이 변경 전에는 기다리는 요청마다 매 스텝 이 로그가 찍혔고, 프롬프트 뒷부분이 Maru 에 없는 요청은 존재 확인도 매 스텝 다시 보냈습니다.
  • _parked_blocks 는 적재 대기 요청마다 ceil(request.num_tokens / block_size) 를 기록합니다. update_state_after_alloc() 에서 적재 대기로 받을 때 넣고, 실행 블록을 받을 때(num_external_tokens <= 0)와 끝나거나 선점될 때(build_connector_meta 정리) 뺍니다.
  • 적재가 실패해 vLLM 이 블록을 돌려받고 다시 묻는 요청은 get_num_new_matched_tokens() 의 처음에서 자기 기록을 지웁니다.
  • _admission_waiters 는 기다리게 한 요청을 먼저 온 순서로 기억하고, 그 몫을 뒤에 온 요청의 예산 계산에 더합니다. 앞에 아무도 없는 요청은 늘 받습니다(vLLM 은 최대 길이 요청 하나가 들어가게 GPU KV 를 잡습니다).
  • 블록 수를 모르면(직접 만든 커넥터 등) 예산을 쓰지 않습니다.
• docs/source/integration/vllm.md 의 비동기 적재 절과 maru_async_load docstring 에 예산을 적었습니다. 이 예산은 GPU 블록을 세고, 기존 maru_load_admission_window 는 동시에 진행하는 적재 수를 셉니다.
• 리뷰 때 봐 주실 점: 예산은 프롬프트 전체 기준이라 GPU 접두사 캐시와 블록을 나눠 쓰는 요청에는 넉넉하게 잡힙니다. 그만큼 적재 대기가 일찍 미뤄질 수 있습니다.

처음 버전에서 뺀 변경

• 처음 올린 버전에는 close() 없이 끝난 핸들러의 영역을 서버가 회수하는 변경도 있었습니다. 실장치에서, 잠시 멈췄을 뿐 아직 살아 있는 핸들러의 영역도 이 변경이 회수해 다른 핸들러에게 내주었고, 멈췄던 핸들러가 깨어나 쓴 데이터가 새 핸들러의 KV 를 덮어썼습니다. 그래서 이 변경을 이 PR 에서 뺐고, 다른 설계로 따로 올립니다.

✅ Tests

• Unit tests
  • 전체 단위 시험(pytest tests/unit, GPU 를 숨기고 MARU_PLUGINS=none 인 환경)은 최종 커밋(b6f20fb)에서 통과 1,110개, 건너뜀 30개입니다. MARU_PLUGINS 를 비우거나 none·임의 이름으로 두어도 플러그인 로더 시험 18개가 모두 통과합니다.
  • 2번 변경(9d9486c)은 시험 코드만 바꿉니다.
  • 3번: tests/unit/test_vllm_connector.py 의 TestDeferredAdmissionFitsCache(11개).
    • 실제 vLLM 스케줄러(make_vllm_scheduler)로 위 배경의 50·25·15블록 시나리오를 돌려, 적재가 끝난 뒤 첫 요청이 실행되는지 확인합니다. 이 변경 바로 앞(9d9486c)의 커넥터로 돌리면 실행된 토큰이 0 이라 실패합니다.
    • 혼자 적재 대기에 들어간 요청도 적재 뒤 실행되는지 확인합니다.
    • 실제 스케줄러에서 예산이 null 블록을 뺀 63블록인지 확인합니다.
    • 커넥터 단위로 예산 초과 시 (None, False), 먼저 기다린 순서 유지, 끝난 요청과 기다리다 끊긴 요청의 몫 해제, 적재 실패 뒤 다시 묻는 요청의 이전 기록 제거, 블록 수를 모를 때의 무제한을 확인합니다.
    • 기다리는 요청을 다시 물어도 자리가 날 때까지 접두사를 조회하지 않고, 자리가 나면 한 번 조회하는지 확인합니다. 조회 결과가 적중 없음으로 바뀐 요청은 자리가 날 때까지 기다린 뒤 재계산으로 넘어가는지도 확인합니다. 두 시험은 b6f20fb 앞의 코드에서 실패합니다.
  • ruff check 와 ruff format --check 를 통과합니다.
• Integration tests
  • 세 변경은 서버와 RPC 코드를 바꾸지 않아 통합 시험은 다시 돌리지 않았습니다. 2번의 autouse 픽스처는 통합 시험에도 적용되지만, 플러그인에 기대는 통합 시험은 없습니다.
• Manual tests

  • 실장치 조건: RTX PRO 6000, vLLM 0.22.1, Qwen2.5-0.5B, --enforce-eager --no-enable-prefix-caching, maru_kv_ops 커널 없는 torch 복사 경로입니다. 비교한 main 은 e84ea9c(이 브랜치가 갈라진 9c0db55 와 문서만 다름)입니다. PR 쪽은 영역 회수 변경을 빼기 전 커밋(ce95214)으로 돌렸습니다. 뺀 변경은 서버·핸들러·RPC 코드만 바꿨고, 세 변경의 커넥터 코드는 b6f20fb 앞까지 지금과 같습니다. b6f20fb 는 실장치에서 다시 돌리지 않았습니다. 측정 호스트는 NVML 이 초기화되지 않아 vLLM 을 NVML 없이 쓰는 CUDA 플랫폼 경로로 띄웠습니다.

  • 1번, vLLM 서빙: cpu 백엔드(풀 8G, 청크 256, --no-async-scheduling)에 매번 새 3,304 토큰 프롬프트를 두 번씩 20회 보냈습니다. TTFT 는 스트리밍 클라이언트가 요청을 보낸 때부터 첫 토큰을 받은 때까지입니다.

    요청	main(20회)	이 PR(20회)
    첫 요청(계산하며 청크 12개 저장)	중앙값 73.6 ms, 최소 53.7, 최대 179.3, 100 ms 초과 6회	중앙값 53.0 ms, 최소 51.1, 최대 58.2, 100 ms 초과 0회
    두 번째 요청(3,072 토큰 Maru 적중)	중앙값 22.3 ms	중앙값 23.1 ms

  • 1번, 커넥터 단독: vLLM 없이 커넥터의 24층 save_kv_layer 를 부르고 앞뒤로 torch.cuda.synchronize() 를 둔 벽시계 시간입니다. 목적지는 익명 mmap(CPU 모드 풀과 같은 종류), 값은 예열 뒤 2회입니다. 이 값은 저장 시간이고, 위 TTFT 와는 다른 양입니다.

    코드	torch CPU 스레드 1개	8개	128개(torch 기본값)
    main	19.5, 19.6 ms	18.8, 18.8 ms	1,823.9, 1,757.9 ms
    이 PR	19.6, 19.7 ms	19.6, 19.8 ms	20.4, 20.7 ms

  • 2번: 장치 플러그인이 설치된 호스트에서 MARU_PLUGINS 를 비운 채 단위 시험 전체를 strace 아래에서 돌렸습니다. 플러그인 모듈 파일은 한 번도 열리지 않았습니다.

  • 3번, vLLM 서빙: 적재하는 엔진을 GPU 블록 64개(블록당 16 토큰), 청크 128 토큰, maru_async_load, CXL-DRAM 장치로 띄웠습니다. 저장하는 엔진이 800 토큰(50블록)·400 토큰(25블록) 프롬프트의 앞 384 토큰씩을 Maru 에 넣은 뒤, 두 프롬프트를 한 completion 요청으로 함께 보냈습니다. main 은 3회 모두 45 s 시간 초과로 멈췄고, 그동안 실행 0·대기 2 였습니다. 이 PR 은 3회에 각각 0.31·0.09·0.09 s 만에 두 요청을 끝냈습니다. 두 요청의 Maru 적중 토큰 합(vLLM 지표 external_prefix_cache_hits 의 증가분)은 3회 모두 768(= 384 + 384)이었습니다.

• No tests needed (reason: )

🔗 Related Issues (optional)

🌿 Related PRs (optional)

• feat(maru_vllm): add CPU and mixed CPU/CXL L1 caching #88 (CPU·mixed 모드를 추가한 PR)

🌿 Related Branches (optional)

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

NEW

CHANGED

FIXED

IMPORTANT NOTES