feat(maru_vllm): remove the LMCache dependency from the direct connector - #75
youngrok-XCENA wants to merge 8 commits into
Conversation
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.
…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.
jooho-XCENA
left a comment
There was a problem hiding this comment.
최신 커밋 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 결과와 비교해 확인했습니다. 구체적인 조건과 수정 방향은 각 인라인 코멘트에 남겼습니다.
| ops.single_layer_kv_transfer( | ||
| layer_dev, | ||
| kv_cache_layer, |
There was a problem hiding this comment.
[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와 호출 인자만 확인하여 이 오류를 잡지 못합니다.
| refusal = self._kv_ops_refusal() | ||
| if refusal is None: | ||
| super().build_extensions() |
There was a problem hiding this comment.
[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 range와 ERROR: Failed building wheel for maru, 종료 코드 1입니다. 비교용으로 TORCH_CUDA_ARCH_LIST=12.0을 지정한 빌드는 성공했습니다.
사용자가 target architecture를 명시하면 GPU 없이도 빌드할 수 있게 유지하되, GPU와 명시적 architecture가 모두 없으면 이 선택적 확장을 생략하도록 처리하고 해당 설치 경로를 검증해 주세요.
🤔 Background & Motivation (Why)
vLLM 직결 커넥터는 LMCache 를 거치지 않고 Maru 에 닿는 것이 존재 이유인데, KV 복사의 기본 경로가
lmcache.c_ops를 부르고 있었습니다. 그 요구사항은 선언된 적도 없습니다.pyproject.toml의 의존성은pyzmq·msgpack·dacite셋뿐입니다.이 의존은 두 가지 방식으로 배포를 해칩니다.
첫째, 없으면 조용히 느려집니다. 오류가 아니라 경고 한 줄을 남기고 레이어별 복사 경로로 내려갑니다. 그 경로의 비용은 커널 채택 당시 실측돼 있습니다. 로드는 동시 1 건 캐시 히트 첫 토큰 148 대 142.4 ms, 저장은 청크당 전송 1 회가 (청크 × 레이어) 당 1 회로 늘어납니다. 배포 입장에서 이것은 오류로 보이지 않고 그냥 느린 장비로 보입니다.
둘째, 기본 경로가 다른 프로젝트의 내부 구조에 매여 있습니다. 커넥터는
ops.EngineKVFormat과ops.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 안에서 처음 드러남"]After — 의존이 사라지고, 커널이 없을 때는 빌드 단계를 지목한 경고가 나갑니다.
flowchart LR C2["MaruKVConnector"] ==> O2["maru_kv_ops<br/>Maru 빌드 산출물"] C2 -.->|"안 빌드됐으면<br/>빌드 단계를 지목한 경고"| F2["레이어별 복사 폴백"] O2 ==> G2["paged KV cache"] F2 --> G2커널 소스는 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"]After — 커널이 slot mapping 을 그대로 받아 주소 계산을 안에서 합니다.
flowchart LR A1["GPU 스테이징 버퍼"] ==>|"slot mapping 직접 소비<br/>블록·칸 계산은 커널 안"| A2["paged KV cache"]바이트 이동과 주소 재배치를 나눈 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__.py—is_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()로 결과를 보고합니다.lmcache.c_ops를 재사용한다」 서술을 제거했습니다. 그 자리에는 Maru 자신의 커널을 만드는 설치 명령과 확인 한 줄이 들어갑니다.커넥터 (
maru_vllm/connector.py)_resolve_kv_ops()— 모듈 수준에서 한 번 해석하고 캐시합니다. 경고는 프로세스당 한 번입니다(로드마다 아님)._packed_load_kernel_ctx의import lmcache.c_ops가 사라지고,self._lmc_ops가self._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가 전처리 검사 실패를 못 막던 문제를 위 설명대로 고쳤습니다.single_layer_kv_transfer가 처리하는 형식 집합이 게이트보다 좁은데 아무 테스트도 그 관계를 고정하지 않던 것을,mem_kernels.cu의 switch 와kv_layout.py의 형식 이름을 읽어 비교하는 테스트로 고정했습니다.test_kv_ops.py가 모듈 최상단importorskip("torch")때문에 CI 에서 통째로 건너뛰어지던 것을 고쳐, 출처·형식 가드가 CI 에서 실제로 돕니다.✅ Tests
단위 테스트 — GPU 를 마스킹한 상태(
CUDA_VISIBLE_DEVICES="")로 확장을 빌드한 호스트에서 937 통과 / 12 skip. 확장이 없는 호스트에서는 927 통과 / 22 skip 이고, 늘어난 skip 은 바인딩 표면 테스트입니다. 커넥터 테스트는 183 → 196,test_kv_ops.py30 개를 새로 추가했습니다.신규 테스트가 고정하는 계약은 여섯입니다.
token_major=False로 넘기는지, 폴백 갈래가num_chunks=1로 넘기는지. 두 인자가 각각 스테이징 모양 해석을 결정합니다.VENDOR.md의 기록과 대조합니다. 새 바이트가 옛 리비전 id 아래 들어오는 것이 실제로 위험한 드리프트인데, 리비전 줄만으로는 잡히지 않습니다.kv_layout.py가 내보내는 형식 네 개가single_layer_kv_transfer의 switch 에 모두 있는지. 게이트는 enum 멤버 존재만 보므로, 갱신이나 새 layout 이 이 관계를 깨면 겹침 경로만 스트림 안에서 예외를 냅니다.컴파일 — CUDA 13.0 / sm_120 으로 컴파일·링크 통과(경고 0). 상류 소스가 LMCache 고유 헤더 없이 단독으로 빌드됩니다. 바인딩 표면이 커넥터가 넘기는 인자 이름과 일치하는 것도 확인했습니다.
설치 경로 — 깨끗한 트리에서 격리 빌드 wheel 에는
maru_kv_ops/_C가 없고--no-build-isolationwheel 에는 있습니다. 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 회) 시간에 따른 장비 상태 변화가 한쪽에만 쌓이지 않게 했습니다./dev/dax1.0--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=0GPU 메모리 비율 0.15 는 동시 8 건의 KV 여유와 스케줄링을 직접 좌우하므로 아래 수치를 읽을 때 함께 보아야 합니다. 절대값은 이 조건의 것이라, 이 본문 앞쪽이 인용한 다른 조건의 첫 토큰 수치(148 대 142.4 ms, 599.7 → 470.2 ms)와는 모델도 프롬프트도 달라 서로 비교할 수 없습니다.
지표의 정체 — 네 가지를 씁니다.
packed-load wall(CXL 읽기 + H2D + 배치 커널 + 스트림 동기화를 감싼 벽시계 시간)을 그 배치가 덮은 요청 수로 나눈 값. 한 배치가 덮는 요청 수가 실행마다 달라서 배치당 값끼리는 비교할 수 없습니다.기본 적재 경로(겹침 끔) — 측정 분해능 안에서 동등
판정 근거는 분해능이 가장 높은 지표입니다. 동시 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_kernel74 개와single_layer_kv_transfer_kernel5 개이고 모두 여기 포함됩니다. 기본 적재 경로에서 이 PR 이 바꾸는 것은ops를 어느 모듈에서 얻는가 하나뿐이며, 그 경로의 나머지 변경분은 주석과 docstring 입니다.겹침 적재 — 발행 비용 절반 감소, 동시 8 건 첫 토큰 −13.3 %
네 지표 모두 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 건입니다. 경로 기록은 구성별로 나뉩니다.
_place_packed_layer()는 커널 갈래와 폴백 갈래 어느 쪽도 로그를 남기지 않으므로, 이 20 회에 대해 커널 갈래를 골랐다는 것 자체는 로그로 확인되지 않습니다. 근거는 기동 시is_available()이 참이라는 것과 발행 시간이 절반으로 줄었다는 것입니다.기능 검증 — 같은 장비에서 아래 구성을 각각 띄웠습니다. 여섯 구성 모두 동시 1 건에서 냉·온 생성 텍스트가 일치했고 오류가 없었습니다.
multi_layer_kv_transfer로 기록됨is_available()참, 겹침 적재가 요청마다 발행 기록을 남김maru_use_layerwise=true경로에 회귀 없음lmcache임포트 차단import lmcache가ImportError를 내도록 막은 채 기동해도 적재·저장 모두 커널 경로로 기록되고 적중 성립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._C는TransferDirection·EngineKVFormat을lmcache.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/VENDOR.md. Built only where PyTorch and nvcc are both present;is_available()reports whether the kernels can be called.CHANGED
maru_vllm/importslmcache, so a change to that package's internal layout can no longer reach the connector's default KV copy path, and thelmcachepackage is not a runtime requirement of the LMCache-free connector. Behaviour and the kernel/fallback gates are unchanged.single_layer_kv_transferinstead 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.--no-build-isolationwhen 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
./install.sh, orpip install -e . --no-build-isolationin an environment with PyTorch and the CUDA toolkit. A plainpip 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.lmcache.c_opscomparison isolates the packaging change. Upstream has since added engine KV formats and refactoredEngineKVFormatonto aKVFormatSpec, 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.lmcache.c_opsbuild: 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 tosingle_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.