Skip to content

feat(maru_vllm): remove the LMCache dependency from the direct connector - #75

Open
youngrok-XCENA wants to merge 8 commits into
mainfrom
feat/vendor-kv-ops
Open

youngrok-XCENA wants to merge 8 commits into
mainfrom
feat/vendor-kv-ops

Conversation

@youngrok-XCENA

@youngrok-XCENA youngrok-XCENA commented Aug 20, 2026

Copy link
Copy Markdown
Collaborator

🤔 Background & Motivation (Why)

vLLM 직결 커넥터는 LMCache 를 거치지 않고 Maru 에 닿는 것이 존재 이유인데, KV 복사의 기본 경로가 lmcache.c_ops 를 부르고 있었습니다. 그 요구사항은 선언된 적도 없습니다. pyproject.toml 의 의존성은 pyzmq·msgpack·dacite 셋뿐입니다.

이 의존은 두 가지 방식으로 배포를 해칩니다.

첫째, 없으면 조용히 느려집니다. 오류가 아니라 경고 한 줄을 남기고 레이어별 복사 경로로 내려갑니다. 그 경로의 비용은 커널 채택 당시 실측돼 있습니다. 로드는 동시 1 건 캐시 히트 첫 토큰 148 대 142.4 ms, 저장은 청크당 전송 1 회가 (청크 × 레이어) 당 1 회로 늘어납니다. 배포 입장에서 이것은 오류로 보이지 않고 그냥 느린 장비로 보입니다.

둘째, 기본 경로가 다른 프로젝트의 내부 구조에 매여 있습니다. 커넥터는 ops.EngineKVFormatops.TransferDirection 처럼 외부 확장의 내부 이름을 다섯 곳에서 직접 읽습니다. 그 이름의 위치와 형태는 Maru 가 정하지도 고정하지도 못하고, 의존이 선언돼 있지 않으니 어느 세대가 설치될지도 알 수 없습니다. getattr(ops.EngineKVFormat, layout.format_name, None) 의 기본값 가드는 형식 이름 하나가 없는 경우만 받아 주고, enum 자체가 그 자리에 없는 경우는 받지 못합니다. 그래서 상류가 내부 배치를 정리하면 그 변화가 커넥터의 기본 KV 복사 경로에서, 그것도 모델 forward 안에서 처음 드러납니다.

이 PR 은 그 묶임을 끊습니다. 병합 후 maru_vllm/ 안에 lmcache 참조가 한 건도 남지 않습니다.

부수적으로 하나 더 있습니다. 레이어 겹침 경로만 KV 배치를 PyTorch 범용 인덱싱으로 하고 있어서, 겹침을 끈 팔(커널 사용)과 켠 팔(PyTorch 사용)의 비교에 배치 구현 차이가 교란으로 섞여 있었습니다.

🏗️ Design Changes

LMCache 의존의 제거

KV 배치 커널을 Maru 안으로 들여옵니다. 커넥터가 의존하는 상대가 외부 패키지에서 Maru 자신의 빌드 산출물로 바뀝니다. 남는 관계는 소스 출처뿐이고, import 도 버전 추적도 없습니다.

Before — 선언되지 않은 외부 의존이 기본 경로에 있습니다. 패키지가 없으면 조용히 느려지고, 있어도 그 내부 이름이 바뀌면 기본 경로가 흔들립니다.

flowchart LR
    C1["MaruKVConnector"] ==>|"import lmcache.c_ops<br/>pyproject 미선언"| K1["LMCache 확장"]
    K1 ==> G1["paged KV cache"]
    C1 -.->|"패키지 없으면<br/>경고 한 줄"| F1["레이어별 복사 폴백"]
    F1 --> G1
    K1 --> X1["내부 이름이 바뀌면<br/>forward 안에서 처음 드러남"]
Loading

After — 의존이 사라지고, 커널이 없을 때는 빌드 단계를 지목한 경고가 나갑니다.

flowchart LR
    C2["MaruKVConnector"] ==> O2["maru_kv_ops<br/>Maru 빌드 산출물"]
    C2 -.->|"안 빌드됐으면<br/>빌드 단계를 지목한 경고"| F2["레이어별 복사 폴백"]
    O2 ==> G2["paged KV cache"]
    F2 --> G2
Loading

커널 소스는 LMCache 에서 바이트 단위로 그대로 가져옵니다(Apache-2.0). 손대지 않는 이유는 갱신을 「파일 복사 + diff」로 유지하기 위해서입니다. 안 쓰는 진입점 다섯 개가 함께 컴파일되지만, 그것을 지우면 「특정 상류 리비전과 동일」이라는 성질이 깨집니다. 바인딩 층은 Maru 자신의 것이고 커넥터가 부르는 진입점만 노출하므로, 상류 서명이 바뀌면 컴파일 오류로 드러납니다. 런타임에 죽는 대신 빌드에서 걸린다는 것이 이 PR 이 바꾸는 핵심입니다.

가져온 리비전은 의도적으로 상류 최신이 아닙니다. 지금까지 성능 측정에 쓰인 LMCache 빌드가 컴파일된 리비전을 고정했습니다. 그래야 이 PR 이 「포장을 바꿨을 뿐」임을 비교로 확인할 수 있습니다. 상류는 그 뒤 배치 형식 두 개를 추가하고 형식 정의를 리팩터해 API 표면을 바꿨는데, 그 채택은 별도 실기 비교가 필요한 작업입니다.

커널은 선택 사항이고, 설치는 그것을 만들어야 한다

Maru 코어는 PyTorch 도 GPU 툴체인도 없는 호스트에 설치되어야 합니다. 그래서 확장은 PyTorch 와 nvcc 가 모두 있을 때만 빌드되고, 없으면 건너뜁니다.

그런데 「선택 사항」이 「기본적으로 안 만들어짐」이 되면 안 됩니다. pip install -e .uv pip install -e . 은 PEP 517 빌드 격리를 기본으로 쓰고, 격리된 환경에는 [build-system].requires 의 setuptools·wheel 만 있습니다. setup.py 는 확장을 torch.utils.cpp_extension 으로 기술하므로, 대상 환경에 PyTorch 가 있어도 빌드를 기술하는 import 자체가 실패해 확장이 생략됩니다. 그리고 그 안내 메시지는 pip 이 삼키는 빌드 서브프로세스 stderr 로 나가서 설치 시점에는 보이지도 않습니다.

그래서 설치 경로가 판단을 대신합니다. install.sh 가 설치 대상 환경에 PyTorch 가 있는지 보고, 있으면 빌드 백엔드를 먼저 넣고 --no-build-isolation 으로 설치한 뒤 커널이 실제로 호출 가능한지 보고합니다. 없으면 격리를 유지하고 어떤 경로로 돌게 되는지 알립니다. 설치 안내 문서에도 같은 내용을 넣었습니다.

빌드가 실패해도 설치는 깨지지 않아야 합니다. 그런데 optional=True 로는 부족합니다. setuptools 는 확장별 컴파일 오류를 흡수하지만, PyTorch 의 build_extensions 는 그 루프 앞에서 CUDA·호스트 컴파일러 전처리 검사를 돌리고, 거기서 나오는 예외는 명령 전체를 타고 올라갑니다. nvcc major 가 PyTorch 가 빌드된 CUDA 와 다르면 pip install 이 그대로 죽습니다. 그 전처리 검사를 미리 돌려서, 거부하면 확장만 목록에서 빼고 설치를 계속합니다.

레이어 겹침의 KV 배치

겹침 경로만 남아 있던 PyTorch 범용 인덱싱을 커널로 바꿉니다. 두 방식은 같은 바이트를 같은 주소에 놓지만, 그 주소를 누가 계산하는지가 다릅니다.

Before — 목적지 주소를 GPU 텐서 두 개로 먼저 만든 뒤 범용 연산에 넘깁니다.

flowchart LR
    B1["GPU 스테이징 버퍼"] --> B2["블록 번호 텐서"]
    B1 --> B3["칸 번호 텐서"]
    B2 --> B4["PyTorch 범용 배치"]
    B3 --> B4
    B4 --> B5["paged KV cache"]
Loading

After — 커널이 slot mapping 을 그대로 받아 주소 계산을 안에서 합니다.

flowchart LR
    A1["GPU 스테이징 버퍼"] ==>|"slot mapping 직접 소비<br/>블록·칸 계산은 커널 안"| A2["paged KV cache"]
Loading

바이트 이동과 주소 재배치를 나눈 2 단 구조는 그대로 둡니다. 커널이 CXL 을 직접 읽으면 매체 읽기 내내 SM 을 붙잡아 같이 도는 디코드를 세우고, 그것을 피하는 것이 비동기 경로가 존재하는 이유입니다(동시 8 건 첫 토큰 599.7 → 470.2 ms, 토큰 간격 25.6 → 21.6 ms 로 확인된 축). 즉 이 PR 이 바꾸는 것은 마지막 배치 단계의 구현 하나입니다.

함께 들여오되 아직 쓰지 않는 것

칸 단위 배치 커널도 같이 가져왔지만 호출하지 않습니다. 자리 지정 단위를 토큰에서 칸으로 바꾸면 metadata 양과 커널 실행 갈래가 16 분의 1 이 되는 후속 축이 있고, 나중에 가져오면 빌드 설정을 두 번 건드려야 합니다. 그 전환의 A/B 판단 기준은 별도 설계 노트에 있습니다.

📝 Implementation Details

패키지 (maru_kv_ops/)

  • csrc/ — 상류에서 그대로 복사한 커널 소스 6 개(2,006 줄). VENDOR.md 에 리비전 40 자 커밋 id, 파일별 SHA-256, 상류가 그 뒤 무엇을 바꿨는지, 갱신 절차를 적었습니다.
  • csrc/pybind.cpp — Maru 자신의 바인딩(90 줄). multi_layer_kv_transfer, single_layer_kv_transfer, multi_layer_block_kv_transfer 셋과 TransferDirection·EngineKVFormat·PageBufferShapeDesc 를 노출합니다. 형식은 12 개 전부 바인딩했습니다. 헤더 enum 이라 비용이 없고, 갱신으로 형식이 늘어도 이 파일을 안 건드립니다.
  • __init__.pyis_available() / import_error(). 확장 import 전에 import torch 가 필요합니다. 확장이 libtorch 에 링크되어 있어 그것이 로더 경로를 잡습니다. 이게 없으면 빌드가 성공해도 libc10.so: cannot open shared object file 로 import 가 실패합니다. 확장은 이름으로 import 합니다. 실행 중인 모듈의 서브모듈을 from-import 하면 미빌드 상태가 "circular import" 로 보고되는데, import_error() 는 운영자에게 그대로 인용되므로 원인을 정확히 말해야 합니다.
  • __getattr__ 로 미빌드 상태의 속성 접근이 빌드 단계를 지목한 AttributeError 를 냅니다.

빌드와 설치 (setup.py, install.sh, 설치 문서)

  • _kv_ops_extension() 이 PyTorch import 가능성과 nvcc 존재를 먼저 확인한 뒤 CUDAExtension 을 만듭니다. nvcc 위치는 PyTorch 에게 물어봅니다. CUDA_HOME 만 설정돼 있고 그 안에 nvcc 가 없으면 PyTorch 의 검사가 예외를 내기 때문에, 게이트와 검사가 같은 것을 봐야 합니다.
  • _optional_kv_ops_build() 가 PyTorch 의 전처리 검사를 빌드 전에 돌리고, 거부하면 확장만 빼고 남은 순수 C 확장을 기본 명령으로 빌드합니다.
  • MARU_SKIP_KV_OPS 로 강제 생략이 가능합니다.
  • install.sh 가 대상 환경의 PyTorch 유무로 --no-build-isolation 을 붙일지 판단하고, 설치 후 is_available() 로 결과를 보고합니다.
  • 설치 문서에 「KV Placement Kernels」 절을 추가했고, vLLM 예제 문서 세 곳에 남아 있던 「LMCache is optional / lmcache.c_ops 를 재사용한다」 서술을 제거했습니다. 그 자리에는 Maru 자신의 커널을 만드는 설치 명령과 확인 한 줄이 들어갑니다.

커넥터 (maru_vllm/connector.py)

  • _resolve_kv_ops() — 모듈 수준에서 한 번 해석하고 캐시합니다. 경고는 프로세스당 한 번입니다(로드마다 아님).
  • _packed_load_kernel_ctximport lmcache.c_ops 가 사라지고, self._lmc_opsself._kv_ops 로 바뀝니다. 게이트(CUDA · Flash 레이아웃 · 형식 존재)는 그대로입니다.
  • _place_packed_layer() 신설 — 커널 갈래와 폴백 갈래를 한자리에 모읍니다. 겹침 경로 두 곳(로더 스레드 발행, 재개된 forward 발행)이 이것을 씁니다.
  • _copy_packed_layer_to_device() 의 스테이징 모양이 [청크, 2, 토큰, hidden][2, 전체 토큰, hidden] 으로 바뀝니다. 커널이 token_major=False 로 읽는 모양이고, 동시에 폴백이 num_chunks=1 로 읽는 모양이라 두 갈래가 한 배치를 공유합니다.
  • _packed_layer_copy_plan() 으로 pitched 복사의 오프셋 산술을 떼어냈습니다. 텐서도 CUDA 도 쓰지 않습니다. 이 부분이 이 PR 에서 가장 위험합니다. 오프셋이 틀리면 옆 레이어의 바이트를 복사하는데 로드는 성공으로 보고합니다. 떼어낸 덕에 테스트가 계획을 합성 슬랩에 재생해 바이트 단위로 검증할 수 있습니다.
  • _packed_store_kernel_ctx 는 이름을 그대로 두고 docstring 을 고쳤습니다. 저장이 첫 사용자였지만 이제 「호출 시점에 레이어 목록이 없는 모든 경로」가 씁니다. 개명은 테스트까지 번져 diff 가 커지므로 분리했습니다. 다만 로그 메시지에서는 방향 표현을 뺐습니다. 이제 로드 경로가 첫 호출자가 될 수 있어서 저장을 한 번도 안 한 시점에 D2H 활성 로그가 찍히기 때문입니다.

리뷰 지적 반영 (커밋 5ca4929, 6ef9f86, 7fae1f6, f2837de, b5fb808)

  • 문서화된 설치 경로가 커널을 만들지 않던 문제와, optional=True 가 전처리 검사 실패를 못 막던 문제를 위 설명대로 고쳤습니다.
  • 벤더 리비전 테스트가 40 자 hex 존재만 확인해서 실제 드리프트를 못 잡던 것을 파일별 SHA-256 비교로 바꿨습니다.
  • single_layer_kv_transfer 가 처리하는 형식 집합이 게이트보다 좁은데 아무 테스트도 그 관계를 고정하지 않던 것을, mem_kernels.cu 의 switch 와 kv_layout.py 의 형식 이름을 읽어 비교하는 테스트로 고정했습니다.
  • test_kv_ops.py 가 모듈 최상단 importorskip("torch") 때문에 CI 에서 통째로 건너뛰어지던 것을 고쳐, 출처·형식 가드가 CI 에서 실제로 돕니다.
  • 여러 청크를 가진 run 이 두 개 이상인 경우의 바이트 검증이 비어 있던 것을 채웠습니다.

✅ Tests

  • Unit tests
  • Integration tests
  • Manual tests

단위 테스트 — GPU 를 마스킹한 상태(CUDA_VISIBLE_DEVICES="")로 확장을 빌드한 호스트에서 937 통과 / 12 skip. 확장이 없는 호스트에서는 927 통과 / 22 skip 이고, 늘어난 skip 은 바인딩 표면 테스트입니다. 커넥터 테스트는 183 → 196, test_kv_ops.py 30 개를 새로 추가했습니다.

신규 테스트가 고정하는 계약은 여섯입니다.

  1. 커널 해석 — 빌드된 확장을 캐시하는지, 미빌드 시 경고가 한 번만 나가고 빌드 단계를 지목하는지, 그리고 미빌드가 폴백으로 이어지고 예외를 내지 않는지.
  2. 오프셋 산술 — 복사 계획을 합성 슬랩에 재생해 「레이어 k 를 요청 토큰 순서로 모은 결과」와 바이트 단위로 비교합니다. 연속 run · 끊어진 run · 여러 청크를 가진 run 두 개 · 단일 청크 run 네 경우를 덮습니다.
  3. 배치 두 갈래 — 커널 갈래가 token_major=False 로 넘기는지, 폴백 갈래가 num_chunks=1 로 넘기는지. 두 인자가 각각 스테이징 모양 해석을 결정합니다.
  4. 출처 무결성 — 벤더 파일 6 개의 SHA-256 을 VENDOR.md 의 기록과 대조합니다. 새 바이트가 옛 리비전 id 아래 들어오는 것이 실제로 위험한 드리프트인데, 리비전 줄만으로는 잡히지 않습니다.
  5. 형식 디스패치kv_layout.py 가 내보내는 형식 네 개가 single_layer_kv_transfer 의 switch 에 모두 있는지. 게이트는 enum 멤버 존재만 보므로, 갱신이나 새 layout 이 이 관계를 깨면 겹침 경로만 스트림 안에서 예외를 냅니다.
  6. 갱신 드리프트 — 진입점·형식·인자 이름이 노출돼 있는지, SPDX 헤더 보존과 리비전 기록이 있는지.

컴파일 — CUDA 13.0 / sm_120 으로 컴파일·링크 통과(경고 0). 상류 소스가 LMCache 고유 헤더 없이 단독으로 빌드됩니다. 바인딩 표면이 커넥터가 넘기는 인자 이름과 일치하는 것도 확인했습니다.

설치 경로 — 깨끗한 트리에서 격리 빌드 wheel 에는 maru_kv_ops/_C 가 없고 --no-build-isolation wheel 에는 있습니다. CUDA 12.8 툴킷과 cu130 PyTorch 조합에서 전처리 검사 거부를 재현했고, 고친 뒤 같은 조건에서 확장만 빠지고 설치는 성공합니다. install.sh 는 PyTorch 있음·없음·확장 미생성 세 갈래를 모두 실행해 확인했습니다.

타입 검사 — 오류 87 건으로 변경 전과 동일(신규 0). 린트·포맷·Sphinx 빌드 통과.

실기 측정 — 기본 적재 경로는 lmcache.c_ops 빌드와 측정 분해능 안에서 같고(동시 8 건 요청당 적재 시간 두 구성 차이 +0.24 %, 같은 구성 안의 실행 간 흔들림 1.3 · 1.7 %), 겹침 적재는 레이어 배치를 발행하는 CPU 시간이 절반으로 줄면서 동시 8 건 첫 토큰이 13.3 % 짧아집니다. 동작 차이는 두 군데뿐입니다. 기본 적재 경로에서 ops 를 어느 모듈에서 얻는가, 그리고 겹침 적재의 배치 구현입니다.

설계 단계에서 기대치를 첫 토큰 개선 0 으로 잡아 두었으므로(전송이 계산보다 173 ms 먼저 끝나 이미 임계경로 밖이고, KV 배치는 레이어당 1 ms 미만), 기본 경로 측정의 목적은 개선 확인이 아니라 회귀 없음 확인입니다. 결과는 그 기대와 맞습니다.

측정 조건 — 이 PR(b5fb808)과 main 병합 지점(84a6aff, lmcache.c_ops 를 쓰는 상태)을 번갈아 5 라운드 돌렸습니다. 한 라운드는 서버를 새로 띄운 독립 실행이고, 홀수 라운드는 이 PR 을 먼저·짝수 라운드는 기준선을 먼저 돌려(3 회 대 2 회) 시간에 따른 장비 상태 변화가 한쪽에만 쌓이지 않게 했습니다.

항목
장비 NVIDIA RTX PRO 6000 Blackwell(sm_120) 1 장, CXL-DRAM /dev/dax1.0
모델 Qwen2.5-0.5B (24 레이어)
프롬프트 동시 요청 슬롯마다 서로 다른 2,725 토큰(청크 256 기준 10 청크)
절차 한 번 저장한 뒤 같은 프롬프트를 되읽기. 동시 1 건은 되읽기 10 회, 동시 8 건은 되읽기 5 회(회마다 8 건을 동시에 발사)
코드 기본값과 다르게 켠 것 --gpu-memory-utilization 0.15, --no-enable-prefix-caching, 응답 길이 32 토큰, maru_pool_size=8G, maru_kv_chunk_tokens=256, maru_log_timing=true, MARU_PLUGINS=none, VLLM_USE_FLASHINFER_SAMPLER=0

GPU 메모리 비율 0.15 는 동시 8 건의 KV 여유와 스케줄링을 직접 좌우하므로 아래 수치를 읽을 때 함께 보아야 합니다. 절대값은 이 조건의 것이라, 이 본문 앞쪽이 인용한 다른 조건의 첫 토큰 수치(148 대 142.4 ms, 599.7 → 470.2 ms)와는 모델도 프롬프트도 달라 서로 비교할 수 없습니다.

지표의 정체 — 네 가지를 씁니다.

  • 요청당 적재 시간: packed-load wall(CXL 읽기 + H2D + 배치 커널 + 스트림 동기화를 감싼 벽시계 시간)을 그 배치가 덮은 요청 수로 나눈 값. 한 배치가 덮는 요청 수가 실행마다 달라서 배치당 값끼리는 비교할 수 없습니다.
  • 레이어 배치 발행 시간: 겹침 경로에서 적재 스레드가 24 개 레이어의 스테이징 복사와 배치를 스트림에 밀어 넣는 데 쓴 CPU 시간을 요청 수로 나눈 값. GPU 실행 시간이 아니라 발행 비용입니다.
  • 두 구성 차이: 라운드마다 각 구성의 중앙값을 내고, 같은 라운드끼리 짝지어 구한 상대 차이. 그 5 개의 중앙값과 범위를 싣습니다. 앞의 ms 두 칸은 라운드 중앙값 5 개의 중앙값이라 ms 두 칸을 나눈 값과 이 칸은 같지 않습니다.
  • 실행 간 흔들림: 한 구성의 라운드별 중앙값 5 개에서 (최대−최소)÷중앙값. 「기준선 · 이 PR」 순으로 두 값을 적으며, 두 구성 차이를 읽을 때의 잡음 바닥입니다.

기본 적재 경로(겹침 끔) — 측정 분해능 안에서 동등

지표 (동시 건수) 이 PR 기준선 두 구성 차이 (중앙값) 차이 범위 실행 간 흔들림 (기준선 · PR)
요청당 적재 시간 (8) 1.37 ms 1.36 ms +0.24 % −0.55 ~ +1.05 % 1.3 % · 1.7 %
요청당 적재 시간 (1) 1.42 ms 1.42 ms −0.34 % −7.9 ~ +5.3 % 9.5 % · 7.7 %
첫 토큰 (8) 50.13 ms 49.84 ms +0.73 % −17.4 ~ +8.9 % 29.0 % · 8.9 %
첫 토큰 (1) 15.37 ms 15.12 ms +0.41 % −3.3 ~ +11.2 % 16.2 % · 10.8 %

판정 근거는 분해능이 가장 높은 지표입니다. 동시 8 건 요청당 적재 시간은 잡음 바닥이 1.3 · 1.7 % 인데 두 구성 차이가 +0.24 % 이고, 라운드별 차이의 부호가 라운드마다 뒤집힙니다. 나머지 세 지표는 차이의 중앙값이 모두 해당 잡음 바닥보다 작습니다. 원래 통과 기준으로 적었던 첫 토큰 차이도 동시 8 건 +0.73 %, 동시 1 건 +0.41 % 로 1 % 안이지만, 첫 토큰은 잡음 바닥이 8.9~29.0 % 라서 그 지표만으로는 1 % 를 가릴 분해능이 없습니다.

이 결과는 기계어 수준에서 설명됩니다. 벤더링한 커널 소스 6 개가 측정 장비에 설치된 LMCache 체크아웃의 같은 파일과 바이트 단위로 동일하고, 두 확장에서 뽑은 sm_120 SASS 를 주소·인코딩 주석을 빼고 대조하면 배치 커널 153 개(명령 39,120 개)가 전부 일치합니다. 그중 72 개는 익명 네임스페이스 인스턴스화라 mangled name 에 번역 단위 해시가 들어가므로 그 해시를 정규화한 뒤 맞췄습니다. 커넥터가 실제로 부르는 것은 load_and_reshape_multi_layer_kernel 74 개와 single_layer_kv_transfer_kernel 5 개이고 모두 여기 포함됩니다. 기본 적재 경로에서 이 PR 이 바꾸는 것은 ops 를 어느 모듈에서 얻는가 하나뿐이며, 그 경로의 나머지 변경분은 주석과 docstring 입니다.

겹침 적재 — 발행 비용 절반 감소, 동시 8 건 첫 토큰 −13.3 %

지표 (동시 건수) 이 PR 기준선 두 구성 차이 (중앙값) 차이 범위 실행 간 흔들림 (기준선 · PR)
레이어 배치 발행 시간 (8) 2.40 ms 5.36 ms −54.6 % −57.7 ~ −51.3 % 9.9 % · 16.9 %
레이어 배치 발행 시간 (1) 0.77 ms 1.51 ms −45.5 % −50.7 ~ −43.2 % 13.3 % · 13.0 %
첫 토큰 (8) 54.41 ms 63.02 ms −13.3 % −19.1 ~ −10.5 % 8.4 % · 8.7 %
첫 토큰 (1) 16.41 ms 17.53 ms −8.6 % −14.6 ~ −4.3 % 9.4 % · 13.4 %

네 지표 모두 5 라운드가 전부 같은 부호입니다. 발행 시간 두 줄은 차이가 잡음 바닥의 세 배를 넘고, 동시 8 건 첫 토큰도 차이 범위 전체가 잡음 바닥 위에 있습니다. 동시 1 건 첫 토큰은 차이가 잡음 바닥과 비슷한 크기라 크기만으로는 약하지만, 5 라운드가 모두 같은 방향입니다.

원인은 발행 횟수입니다. 기준선은 레이어마다 목적지 주소를 GPU 텐서 두 개로 만든 뒤 PyTorch 범용 연산에 넘겨 레이어당 커널 3 회를 발행하는데, 이 PR 은 커널 1 회로 끝냅니다. 24 레이어 기준으로 요청당 발행이 72 회에서 24 회로 줄고, 그것이 발행 시간 절반 감소로 나타납니다. 다만 발행 시간은 적재 스레드의 CPU 비용이므로 그 감소분이 곧 첫 토큰 단축은 아닙니다. 첫 토큰 단축은 따로 측정한 결과이고, 발행 비용 감소는 그 메커니즘입니다.

이 표는 겹침을 켠 두 구현을 비교한 것이지 겹침 켬·끔 비교가 아닙니다. 이 조건에서 겹침 켬은 이 PR 안에서도 끔보다 느립니다. 같은 라운드끼리 짝지으면 동시 8 건 첫 토큰이 5 라운드 모두 겹침 쪽이 크고(라운드별 +1.7 ~ +11.8 %, 중앙값 +10.8 %), 동시 1 건은 5 라운드 중 4 개가 그렇습니다(중앙값 +11.4 %). 겹침의 순이득은 계산이 전송보다 큰 조건에서 나오므로 그 판정은 이 측정의 범위 밖입니다.

기능 게이트 — 측정 40 회 실행 전부에서 저장 실패 0 건, 적재 미스 0 건, 재계산 0 건입니다. 경로 기록은 구성별로 나뉩니다.

  • 기본 적재 경로 20 회: 적재 커널 호출이 동시 1 건 10 회·동시 8 건 40 회로 되읽기 요청 수와 정확히 같고, 폴백 기록은 없습니다. 두 구성이 동일합니다.
  • 겹침 경로 20 회: 적재 스레드의 레이어 배치 발행 기록이 동시 1 건 10 회·동시 8 건 40 회로 요청 수와 같습니다. 두 구성이 동일합니다. 다만 _place_packed_layer() 는 커널 갈래와 폴백 갈래 어느 쪽도 로그를 남기지 않으므로, 이 20 회에 대해 커널 갈래를 골랐다는 것 자체는 로그로 확인되지 않습니다. 근거는 기동 시 is_available() 이 참이라는 것과 발행 시간이 절반으로 줄었다는 것입니다.

기능 검증 — 같은 장비에서 아래 구성을 각각 띄웠습니다. 여섯 구성 모두 동시 1 건에서 냉·온 생성 텍스트가 일치했고 오류가 없었습니다.

구성 확인한 것
기본 적재 적재·저장이 multi_layer_kv_transfer 로 기록됨
겹침 켬 기동 시 is_available() 참, 겹침 적재가 요청마다 발행 기록을 남김
레이어별 저장 maru_use_layerwise=true 경로에 회귀 없음
인스턴스 2 대 공유 한쪽이 저장하고 다른 쪽이 적재, 양쪽 모두 커널 경로로 기록됨
lmcache 임포트 차단 import lmcacheImportError 를 내도록 막은 채 기동해도 적재·저장 모두 커널 경로로 기록되고 적중 성립
커널 부재 is_available() 이 거짓일 때 예외 없이 레이어별 폴백으로 내려가고, 경고가 빌드 단계를 지목

동시 8 건에서는 냉·온 텍스트가 8 개 요청 중 4 개에서 갈립니다. 같은 갈림이 병합 전 main 에서도 같은 4 개 요청·같은 문자 위치(56, 54, 136, 56)에서 나오고, Maru 를 끄고 vLLM 자체 prefix cache 만 쓴 대조 실행에서도 그중 3 곳이 같은 위치에서 갈립니다. 즉 동시 실행 시의 배치 비결정성이며 이 PR 이 만든 차이가 아닙니다. 텍스트 바이트 일치는 동시 실행에서 깨지는 판정 기준이라 단독 근거로 쓰지 않았습니다.

문서대로 설치한 환경에서의 기동도 확인했습니다. CUDA 13.0 툴체인으로 sm_90·sm_120 확장이 경고 없이 빌드되고 is_available() 이 참입니다.

남는 한계 — 커널 경로와 레이어별 폴백을 런타임에 갈라 볼 설정 손잡이는 없습니다. 커널 경로는 쓸 수 있으면 무조건 기본이고, MARU_SKIP_KV_OPS 는 설치 시점 환경변수라 서빙 중 A/B 에 쓸 수 없습니다. 폴백 경로를 측정 대상으로 삼을 일이 생기면 별도 PR 에서 손잡이를 추가하는 것이 맞습니다.

병합 전 필요 항목으로 적었던 세 건 중 두 건(내부화 전후 비교, 문서대로 설치한 환경에서의 기동)은 위에 측정 결과가 있습니다. 나머지 한 건인 겹침 켬·끔 이득 재확정은 이 PR 범위 밖으로 옮깁니다. 위 측정이 보여주듯 그 판정은 계산이 전송보다 큰 조건을 골라야 성립하고, 겹침 이득의 대외 인용값을 확정하는 별도 후속과 같은 작업이기 때문입니다. 이 PR 이 책임지는 범위는 겹침을 켠 상태에서 배치 구현을 바꾼 것이 회귀가 아니라는 확인까지입니다.

maru_kv_ops._CTransferDirection·EngineKVFormatlmcache.c_ops 와 같은 pybind11 타입 이름으로 등록합니다. 그래서 한 프로세스가 둘 다 임포트하면 나중에 임포트한 쪽이 실패합니다. 측정 장비에서 양방향으로 재현했습니다. lmcache.c_ops 가 먼저 들어온 프로세스에서는 maru_kv_ops.is_available() 이 거짓이 되어 커넥터가 폴백으로 내려가고, 반대 순서에서는 LMCache 가 Failed to import backend lmcache.c_ops: generic_type: type "TransferDirection" is already registered! 를 남기고 자기 백엔드를 못 씁니다. lmcache 가 깔리지 않은 직결 커넥터 배포에서는 발생하지 않지만, 두 패키지가 함께 깔린 호스트에서는 임포트 순서 하나로 경로가 갈립니다. 두 enum 에 py::module_local() 을 붙이면 해결되므로 별도로 다루는 편이 좋겠습니다.

🔗 Related Issues (optional)

🌿 Related PRs (optional)

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

NEW

  • maru_kv_ops: new package holding Maru's paged-KV placement kernels, vendored from LMCache under Apache-2.0 with the source revision, a per-file SHA-256 and the refresh procedure recorded in maru_kv_ops/VENDOR.md. Built only where PyTorch and nvcc are both present; is_available() reports whether the kernels can be called.

CHANGED

  • maru_vllm: the direct connector no longer depends on LMCache in any way. Nothing under maru_vllm/ imports lmcache, so a change to that package's internal layout can no longer reach the connector's default KV copy path, and the lmcache package is not a runtime requirement of the LMCache-free connector. Behaviour and the kernel/fallback gates are unchanged.
  • maru_vllm: the layer-overlap load places KV with single_layer_kv_transfer instead of PyTorch advanced indexing, which drops two index tensors and two kernel launches per layer and removes a scatter-implementation difference from every overlap on/off comparison. The staging tensor becomes [2, num_tokens, hidden], which both the kernel and the per-layer fallback read directly.
  • maru_vllm: an unavailable placement kernel now logs once per process naming the build step that fixes it, rather than a single line stating the import failed.
  • install.sh installs with --no-build-isolation when the target environment has PyTorch, so the documented install builds the placement kernels instead of silently skipping them, and reports afterwards whether they are callable. A CUDA toolkit that PyTorch's preflight rejects now drops the extension instead of failing the install.

FIXED

IMPORTANT NOTES

  • The kernels are a build product of this package, so an existing install has to be redone to get them: run ./install.sh, or pip install -e . --no-build-isolation in an environment with PyTorch and the CUDA toolkit. A plain pip install -e . runs an isolated build that never sees PyTorch and skips them. Without them the connector runs its per-layer fallback, which measured 148 vs 142.4 ms on a single-request cache-hit load and turns a chunk's single store transfer into one per (chunk, layer). python -c "import maru_kv_ops; print(maru_kv_ops.is_available())" reports which path a deployment will take.
  • The vendored revision is deliberately not upstream's latest: it is the revision the measured LMCache build was compiled from, so a vendored-vs-lmcache.c_ops comparison isolates the packaging change. Upstream has since added engine KV formats and refactored EngineKVFormat onto a KVFormatSpec, changing the API surface; adopting that needs its own on-device comparison. Because the sources are vendored, such a change is now a compile error at refresh time rather than a runtime crash.
  • On-device measurement, five rounds alternating this PR with the pre-merge main, puts the default load path inside measurement resolution of the lmcache.c_ops build: per-request load time differs by +0.24% at concurrency 8, where the run-to-run spread within either configuration is 1.3% and 1.7%, and the per-round difference changes sign from round to round. The vendored sources compile to byte-identical sm_120 SASS across all 153 placement kernels, so that is expected rather than lucky. The layer-overlap load, whose placement changed from PyTorch advanced indexing to single_layer_kv_transfer, issues a layer's placement in about half the CPU time and lowers first-token latency by 13.3% at concurrency 8, with all five rounds between -10.5% and -19.1%. Cache hits, store failures, load misses and recomputes were identical across both configurations. Reconfirming the overlap on/off gain is left out of this PR; see the Tests section.

The vLLM connector reached Maru without LMCache, yet its default KV copy
path called lmcache.c_ops. That made LMCache a runtime requirement of the
connector whose purpose is not to need one, and the requirement was never
declared: pyproject listed pyzmq, msgpack and dacite, and a missing
LMCache degraded to a per-layer copy behind a single log line — 148 vs
142.4 ms on a single-request cache hit, and a store transfer per
(chunk, layer) instead of one per chunk. A deployment reads that as a
slower host, not as a misconfiguration.

Vendor the kernels into maru_kv_ops (Apache-2.0, revision and refresh
procedure in maru_kv_ops/VENDOR.md) and drop the lmcache import. The
sources are copied byte-for-byte so a refresh stays a file copy and a
diff; pybind.cpp is ours and binds only the three entry points the
connectors call, so an upstream signature change surfaces as a compile
error. The extension is built only where PyTorch and nvcc are present,
because Maru's core installs on hosts with neither, and it is the runtime
that is made loud instead: an unbuilt extension logs once, naming the
build step, rather than falling back in silence.

The revision pinned is the one the measured LMCache build was compiled
from, not upstream's latest. Upstream has since added engine KV formats
and refactored EngineKVFormat onto a KVFormatSpec, which changes the API
surface; adopting that needs its own on-device comparison, and pinning
the measured revision is what lets this change be compared against
lmcache.c_ops as a packaging change alone.

Also route the layer-overlap load through single_layer_kv_transfer. It
was the one copy path still placing KV with PyTorch advanced indexing,
which cost two index tensors and three kernel launches per layer where
the kernel needs one, and — because the overlap-off arm already used the
kernel — put a scatter-implementation difference inside every overlap
on/off comparison. The staging tensor becomes [2, num_tokens, hidden] to
match what the kernel reads, which is also the shape the per-layer
fallback reads as a single run. The copy-engine-then-kernel split is
unchanged: a kernel reading CXL directly holds SMs for the whole media
read and stalls co-scheduled decode, which is what the deferred path
exists to avoid.

The pitched-copy offsets move into _packed_layer_copy_plan, free of
tensors and CUDA, because a wrong offset there copies the neighbouring
layer's bytes and still reports success. Tests replay the plan against a
synthetic slab and compare it with the gather it stands for.

Tests: 919 unit tests pass with the GPU masked off; the extension
compiles for sm_120 and its binding surface is asserted against the
argument names the connector passes. No on-device measurement was run,
so the equivalence and overlap claims above remain to be measured.
@github-actions

github-actions Bot commented Aug 20, 2026

Copy link
Copy Markdown

youngrok-XCENA and others added 6 commits September 2, 2026 21:43
…tall

The kernels this package now owns were never built by the install the docs
tell people to run. `install.sh` and a hand-rolled `pip install -e .` both use
PEP 517 build isolation, whose environment holds only
`[build-system].requires` — setuptools and wheel. `setup.py` describes the
extension with `torch.utils.cpp_extension`, so the import that decides whether
to build it fails in that environment no matter what the target environment
has, and the extension is skipped. pip captures the build's stderr, so the
skip notice never reaches the operator either. Verified on a clean tree: an
isolated `pip wheel` ships `maru_shm/_cxl_flush` and no `maru_kv_ops/_C`,
while the same command with `--no-build-isolation` ships both.

install.sh now looks for PyTorch in the environment it is installing into,
adds the build backend and `--no-build-isolation` when it finds it, and reports
afterwards whether the kernels are callable. Without PyTorch it keeps isolation
and says which path the connector will take. The installation guide gains the
same for a by-hand install.

`optional=True` also does not hold the promise it was given. setuptools
absorbs a per-extension compile error, but PyTorch's `build_extensions` opens
with a CUDA and host-compiler preflight that raises out of the whole command,
so an nvcc major differing from the one PyTorch was built against takes
`pip install` down with it. Reproduced on a host with CUDA 12.8 and a cu130
PyTorch: `RuntimeError: The detected CUDA version (12.8) mismatches ...`, exit
1. The preflight is now run before the build and a refusal drops the extension
instead, leaving the plain C extension to build and the install to succeed.
The nvcc gate asks PyTorch where the toolkit is rather than trusting
`CUDA_HOME` to contain one, which is its own way to reach that same failure.

`_IMPORT_ERROR` is quoted verbatim to operators by both the connector warning
and the new install report, so the extension is imported by name: a
from-import of a submodule of the module still executing reports the partially
initialised package and suggests a circular import, where the actual cause is
that nothing was compiled.
Three gaps in what the vendoring was pinned against.

The revision test only asserted that some 40-character hex token appears in
VENDOR.md. A refresh that copies new files and leaves the table alone passes
it, which is exactly the silent drift the test says it prevents. VENDOR.md now
records a SHA-256 per copied file and the test recomputes them, so the record
is a claim the suite checks. The digests are those of the recorded revision:
verified against `43a3318` in a LMCache checkout, all six identical.

Nothing tied the formats the connector can hand a kernel to the formats the
kernel dispatches. The gate accepts any name the `EngineKVFormat` enum
carries, because it was written for `multi_layer_kv_transfer`; the
layer-overlap load calls `single_layer_kv_transfer`, whose switch covers a
narrower set and raises on the rest, inside a CUDA stream at serving time. The
two agree today, and a test now reads the switch out of `mem_kernels.cu` and
the format names out of `kv_layout.py` so a refresh or a new layout has to
keep them agreeing.

The module opened with `pytest.importorskip("torch")`, and the `dev` extra CI
installs carries no torch, so the whole file was skipped in the one place
these guards could fire before a review. Nothing outside the binding-surface
tests needs torch — the package's own import wraps it — so the skip now sits
on that class alone.

Checked with torch masked out of `sys.meta_path`: 1 skipped before, 19 passed
after. Both new guards were confirmed to fail on a mutated source tree.
The copy plan was covered byte-for-byte in two shapes: one run spanning every
chunk, and two runs of a single chunk each. Chunked prefill's actual shape —
several runs of several chunks — was covered only by a count of copies and
rows, and it is the shape where the destination arithmetic can be wrong
without a symptom, because a run has to resume inside the K half where the
previous run stopped rather than at the half's start.

The new case replays a plan for two two-chunk runs against a synthetic slab
with a gap between them, and compares the result with the gather it stands
for. Confirmed to be the only test that catches a `chunk_start` mistaken for
`run_index`: with that substitution the existing cases pass, since their run
indices and chunk starts coincide.
…ck notes

Three comments that describe something other than what the code does.

The staged run views are a function-local list, so they cannot be what keeps
the CXL mapping alive across queued copies as the comment claimed; the callers
retaining `infos` against the stream's completion event are. Stating the wrong
guarantee is worse than stating none, because the next reader trusts it.

`_packed_store_kernel_ctx` now resolves on whichever path gets there first,
and the asynchronous loads do, so a run that has not stored anything yet logs
"coalesced kernel D2H enabled". The message no longer names the store or a
direction. Its docstring also still said an unusable kernel is cached as
`False`, where the code sets a `_store_kernel_unusable` flag.

`_place_packed_layer` claims both branches place the same bytes at the same
addresses. That holds for the layout branch of `_inject_kv_into_layer`, which
is what the new `[2, num_tokens, hidden]` staging tensor was checked against;
its MLA and legacy rank-4 Triton branches read the buffer as token-major
regardless of `num_chunks`. They derive the same shapes from this tensor as
from the chunk-major one, and those shapes do not fit the paged tensor they
assign to, so the outcome is the raise and recompute it already was — which is
worth saying next to a claim of equivalence rather than leaving to be
rediscovered.
…ctor

The connector's reason to exist is reaching Maru without LMCache, and the
kernels it needs are now Maru's own. Three user-facing pages still told
readers the opposite — that "LMCache is optional", and that installing its
Python package is what gets them the coalesced multi-layer transfers. After
this branch nothing in `maru_vllm/` imports `lmcache`, so those paragraphs
describe a coupling that no longer exists and would send a reader off to
install a package that has no part in these examples.

They are removed rather than reworded. A page that keeps explaining the
relationship still frames the connector in terms of LMCache; the requirement
the reader actually has is that Maru's own kernels be built.

That requirement takes their place, because the same pages were also telling
readers to install with a command that never builds them: `uv pip install -e
/path/to/maru` runs an isolated build with no PyTorch in it. They now point at
install.sh or the same command with `--no-build-isolation`, and give the
one-liner that reports which path the connector will take.

Two test docstrings named `c_ops` as the kernel they mirror; they now name
maru_kv_ops. The remaining LMCache references in docs and examples belong to
`maru_lmcache`, the opposite integration where Maru is LMCache's storage
backend, and are left alone.
@youngrok-XCENA youngrok-XCENA changed the title feat(maru_kv_ops): own the paged-KV placement kernels feat(maru_vllm): remove the LMCache dependency from the direct connector Sep 3, 2026
@youngrok-XCENA
youngrok-XCENA marked this pull request as ready for review September 15, 2026 00:54
@youngrok-XCENA
youngrok-XCENA requested a review from a team September 15, 2026 05:13

@jooho-XCENA jooho-XCENA left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

최신 커밋 1508da3 기준으로 확인했고, 아래 두 문제를 재현하여 수정 요청드립니다.

  • P1: HND 레이아웃에서 새 single-layer 커널이 블록 크기와 헤드 수를 뒤집어 해석하여 KV 데이터를 잘못 기록합니다.
  • P2: PyTorch와 nvcc는 있지만 GPU가 보이지 않는 빌드 환경에서 선택적 확장을 생략하지 못하고 패키지 빌드 전체가 실패합니다.

검증 환경은 PyTorch 2.11.0+cu128 / CUDA toolkit 12.8입니다. TORCH_CUDA_ARCH_LIST=12.0을 지정한 확장 컴파일은 성공했습니다. GPU를 마스킹한 전체 단위 테스트는 992 passed / 9 skipped / 147 deselected였고, HND 오류는 실제 GPU에서 기존 fallback 결과와 비교해 확인했습니다. 구체적인 조건과 수정 방향은 각 인라인 코멘트에 남겼습니다.

Comment thread maru_vllm/connector.py
Comment on lines +2353 to +2355
ops.single_layer_kv_transfer(
layer_dev,
kv_cache_layer,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] HND 텐서를 커널이 기대하는 물리적 축 순서의 view로 전달해야 합니다.

vLLM의 HND 캐시는 물리적으로 HND 순서여도 등록된 텐서의 shape은 [NB, 2, BS, NH, HS] 또는 [2, NB, BS, NH, HS]입니다(kv_layout.py에도 이 계약이 설명되어 있습니다). 반면 이번에 호출하는 single_layer_kv_transfer의 HND 분기는 size(2)num_heads, size(3)block_size로 읽습니다. 여기서 kv_cache_layer를 그대로 넘기면 BS와 NH가 뒤바뀌어, overlap 로드가 잘못된 KV 주소에 기록됩니다. 기존 _inject_kv_into_layer는 논리적 shape과 stride에 따라 접근하므로 정상입니다.

실제 GPU에서 NB=64, BS=16, NH=8, HS=16, 16개 토큰의 float32 데이터를 사용하여 기존 fallback 결과와 비교했습니다. NHD 두 형식은 일치했지만 HND 두 형식(NL_X_NB_TWO_NH_BS_HS, NL_X_TWO_NB_NH_BS_HS)은 각각 7,936개 원소가 불일치했습니다. 같은 테스트에서 HND 목적지를 kv_cache_layer.permute(0, 1, 3, 2, 4) view로 전달하면 불일치가 0이 됩니다. 재현 시 잘못된 source read가 할당 범위를 벗어나지 않도록 staging의 backing allocation에 여유 공간을 두었습니다.

HND 분기에서 물리적 축 순서의 view를 전달하거나 커널이 실제 layout 정보를 받도록 수정하고, BS != NH인 HND에서 kernel/fallback의 실제 출력 데이터를 비교하는 회귀 테스트를 추가해 주세요. 현재 테스트는 format dispatch와 호출 인자만 확인하여 이 오류를 잡지 못합니다.

Comment thread setup.py
Comment on lines +117 to +119
refusal = self._kv_ops_refusal()
if refusal is None:
super().build_extensions()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] GPU가 보이지 않고 target architecture가 미설정인 빌드도 선택적 확장 생략으로 처리해야 합니다.

PyTorch와 nvcc가 설치되어 있지만 GPU가 보이지 않는 컨테이너/빌드 호스트에서는 _kv_ops_refusal()의 CUDA/ABI 검사가 통과합니다. 이후 super().build_extensions()가 아키텍처를 자동 감지하면서 PyTorch의 _get_cuda_arch_flags() 안에서 빈 arch_list[-1]에 접근하여 IndexError를 냅니다. 이 예외는 optional=True로 흡수되지 않아 확장만 빠지는 대신 Maru 설치 전체가 실패합니다. install.sh가 PyTorch를 발견하면 자동으로 비격리 빌드를 선택하므로 문서화된 설치 경로에서도 영향을 받습니다.

PyTorch 2.11.0+cu128 / nvcc 12.8 환경의 깨끗한 소스에서 다음 명령으로 재현했습니다:

env -u TORCH_CUDA_ARCH_LIST CUDA_VISIBLE_DEVICES='' \
  python -m pip wheel --no-deps --no-build-isolation .

결과는 IndexError: list index out of rangeERROR: Failed building wheel for maru, 종료 코드 1입니다. 비교용으로 TORCH_CUDA_ARCH_LIST=12.0을 지정한 빌드는 성공했습니다.

사용자가 target architecture를 명시하면 GPU 없이도 빌드할 수 있게 유지하되, GPU와 명시적 architecture가 모두 없으면 이 선택적 확장을 생략하도록 처리하고 해당 설치 경로를 검증해 주세요.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants