feat(maru_vllm): place the layer-overlap load with single_layer_kv_transfer instead of PyTorch indexing

[youngrok-XCENA]

배경

CXL 에서 GPU 로 KV 를 가져오는 일은 두 단계입니다. 바이트를 스테이징 버퍼로 옮기고, 그 내용을 paged KV cache 의 제자리(블록 번호 + 칸 번호)에 흩뿌립니다. 두 번째 단계를 누가 하느냐가 경로마다 다릅니다.

경로	두 번째 단계
기본 적재 (겹침 끔)	multi_layer_kv_transfer 커널 한 번이 slot mapping 을 받아 주소를 안에서 계산
겹침 적재 (maru_overlap_load_with_compute)	PyTorch 범용 인덱싱 — 블록 번호 텐서와 칸 번호 텐서를 먼저 만든 뒤 index_put

겹침 경로는 레이어마다 목적지 주소를 GPU 텐서 두 개로 만들고 커널을 여러 번 발행합니다. 이것을 기본 경로처럼 커널 한 번으로 바꾸면 24 레이어 기준 요청당 발행 횟수가 줄어듭니다.

maru_kv_ops 는 이미 그 진입점(single_layer_kv_transfer)을 바인딩해 두었습니다(#85). 다만 현재 아무도 부르지 않습니다.

하려는 것

겹침 적재의 배치 단계를 single_layer_kv_transfer 로 옮깁니다.

왜 지금 당장은 아닌가

이득이 이 구조 위에서 측정된 적이 없습니다. 닫힌 PR #75 에서 잰 「동시 8 건 첫 토큰 −13.3 %」는 #78 이전 겹침 구현과 비교한 값입니다. #78 이 겹침을 chunkwise·layerwise 두 갈래로 재구성했으므로 그 수치는 현재 코드에 적용되지 않습니다. 착수 전에 현재 구조에서 이득이 있는지부터 재야 합니다.

설계가 필요합니다. #78 의 _copy_layerwise_layer_to_device 가 만드는 평탄 텐서는 single_layer_kv_transfer 가 읽는 모양이 아닙니다. 두 갈래 중 chunkwise 만 커널로 보낼지, 커널 입력 모양을 맞춰 둘 다 보낼지를 정해야 합니다.

반드시 안고 가야 할 것 — HND 배치의 조용한 손상

#75 에서 이 변경을 시도했을 때 HND 캐시에서 KV 가 잘못된 주소에 기록됐습니다. 원인과 재현, 수정, 회귀 테스트가 거기 남아 있으니 그대로 가져가야 합니다.

single_layer_kv_transfer 의 HND 분기는 목적지 텐서의 size(2) 를 헤드 수, size(3) 을 페이지 크기로 읽습니다. 즉 물리 축 순서가 shape 에 드러나 있다고 가정합니다. 그런데 vLLM 은 HND 로 할당한 뒤 곧바로 되돌려 permute 하므로, 등록된 텐서는 백엔드의 토큰 우선 shape 을 그대로 지니고 HND 는 stride 에만 남습니다(maru_vllm/kv_layout.py 의 KVLayout docstring 에 이 계약이 적혀 있습니다). 그래서 두 값이 뒤바뀝니다.

기존 PyTorch 인덱싱 경로는 논리 shape 으로 색인하고 stride 를 따라가므로 영향이 없습니다. 커널로 옮기는 순간 생기는 회귀입니다.

커널 수준 재현 (NB=64, BS=16, NH=8, HS=16, 16 토큰 float32, 커널 출력 대 레이어별 폴백 출력):

커널 형식	불일치 원소
NL_X_TWO_NB_BS_NH_HS (NHD)	0 / 262,144
NL_X_NB_TWO_BS_NH_HS (NHD)	0 / 262,144
NL_X_TWO_NB_NH_BS_HS (HND)	9,536 / 262,144
NL_X_NB_TWO_NH_BS_HS (HND)	7,168 / 262,144

엔드투엔드로는 VLLM_KV_CACHE_LAYOUT=HND + 겹침 켬에서 캐시 적중을 5.46 배 개선으로 보고하면서 요청과 무관한 글을 생성했습니다. 오류도 경고도 남지 않습니다.

첫 실행: The three components of a KV cache sharing system are chunk-based storage, ...
재실행 : Condorcets are also known as the 5-C F595, Pogansson, 4.0 and 5G-600000002-003 ...

수정 — HND 일 때 목적지를 물리 축 순서 view 로 넘깁니다. 위 표의 불일치가 전부 0 이 됩니다.

회귀 테스트 — rank-5 네 형식 각각에서 커널 출력과 레이어별 폴백 출력을 바이트 단위로 대조합니다. BS != NH 로 잡아야 합니다. 두 값이 같으면 축이 뒤바뀌어도 결과가 같아 결함이 숨습니다.

수정과 테스트는 닫힌 브랜치 feat/vendor-kv-ops 의 커밋 12255bb 에 있습니다.

착수 조건

1. 현재 구조(feat(maru_vllm): run the load/compute overlap on the layerwise storage format #78 의 두 갈래)에서 겹침 배치를 커널로 바꿨을 때의 이득을 먼저 측정한다. 이득이 잡음 바닥 안이면 하지 않는다.
2. 하기로 하면 위 HND 수정과 회귀 테스트를 함께 넣는다.
3. 검증 격자에 캐시 축 순서(NHD·HND) 와 저장 형식(chunkwise·layerwise) 을 둘 다 넣는다. feat(maru_vllm): remove the LMCache dependency from the direct connector #75 의 검증이 NHD 만 돌아 이 결함을 놓쳤습니다.

관련

• feat(maru_vllm): own the paged-KV placement kernels #85 — 배치 커널을 Maru 안으로 들여온 PR. single_layer_kv_transfer 를 바인딩하지만 호출하지는 않습니다. 이 이슈가 그 첫 호출자가 됩니다.
• feat(maru_vllm): run the load/compute overlap on the layerwise storage format #78 — 겹침을 레이어별 저장 형식에서도 돌게 만든 PR. 이 이슈가 올라탈 구조입니다.
• feat(maru_vllm): remove the LMCache dependency from the direct connector #75 (closed) — 이 변경을 feat(maru_vllm): own the paged-KV placement kernels #85 와 함께 시도했다가 분리한 PR. HND 결함과 수정, 회귀 테스트가 거기 있습니다.
• feat(maru_vllm): maru-native CUDA kernel for packed GPU↔CXL KV transfers #74 — Maru 가 직접 작성한 stride 기반 커널. 이 이슈는 기존 커널을 더 많은 경로에서 부르는 일이고, feat(maru_vllm): maru-native CUDA kernel for packed GPU↔CXL KV transfers #74 는 커널 자체를 새로 쓰는 일이라 별개입니다.