LLM 스터디 3주차 - PagedAttention 직접 증명하기 — vLLM 블록 테이블 꺼내 보기

python3 -m venv ~/pa-lab
~/pa-lab/bin/pip install vllm

os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"

llm.llm_engine.engine_core.engine_core.scheduler
   .kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks

내가 준 이름   : R1
엔진이 붙인 이름 : R1-bf6e0edb

def blocks_of(tag):
    for request_id, blocks in manager.req_to_blocks.items():
        if request_id == tag or request_id.startswith(tag + "-"):
            return [b.block_id for b in blocks]
    return None

llm = LLM(model="HuggingFaceTB/SmolLM-135M",
          max_model_len=512, block_size=16,
          num_gpu_blocks_override=80,      # 풀을 80블록으로 좁힌다
          enable_prefix_caching=False,     # 나눠 쓰기를 꺼야 구멍만 보인다
          gpu_memory_utilization=0.35, enforce_eager=True)

[67, 68, ... 79,   44, 43, ... 31]
 └─ 올라간다 ─┘   └─ 내려간다 ─┘

"""PagedAttention 검증 1 — 블록 단위 할당과 prefix 공유

vLLM 엔진을 같은 프로세스 안에서 띄우고, 스케줄러가 들고 있는
블록 테이블(req_to_blocks)을 직접 꺼내 본다.
"""
import os

# 엔진을 별도 프로세스가 아니라 이 프로세스 안에서 돌린다. 내부를 보려면 필수.
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
# nvcc 가 없는 환경에서 FlashInfer 샘플러가 JIT 컴파일하다 죽는 것을 피한다.
os.environ["VLLM_USE_FLASHINFER_SAMPLER"] = "0"
os.environ["TORCHDYNAMO_DISABLE"] = "1"

import math
import torch
from vllm import LLM, SamplingParams

BLK = 16                      # 블록 하나가 담는 토큰 수
MAX_LEN = 512                 # 최대 컨텍스트 길이
BAR = "=" * 74

llm = LLM(model="HuggingFaceTB/SmolLM-135M",
          max_model_len=MAX_LEN,
          block_size=BLK,
          enable_prefix_caching=True,
          gpu_memory_utilization=0.35,
          enforce_eager=True)

engine = llm.llm_engine
scheduler = engine.engine_core.engine_core.scheduler
manager = scheduler.kv_cache_manager.coordinator.single_type_managers[0]
tokenizer = engine.get_tokenizer()
cache_config = engine.vllm_config.cache_config


def blocks_of(tag):
    """요청이 들고 있는 물리 블록 번호 목록.

    엔진이 request_id 뒤에 접미사를 붙이므로("R1" -> "R1-bf6e0edb")
    정확히 일치시키지 말고 접두사로 찾는다.
    """
    for request_id, blocks in manager.req_to_blocks.items():
        if request_id == tag or request_id.startswith(tag + "-"):
            return [b.block_id for b in blocks]
    return None


def add(tag, prompt, max_tokens=4):
    engine.add_request(tag, prompt,
                       SamplingParams(temperature=0, max_tokens=max_tokens))


def drain():
    """남은 요청을 모두 끝낸다."""
    while engine.has_unfinished_requests():
        engine.step()


# ─────────────────────────────────────────────────────────────
print("\n" + BAR)
print("[1] KV 캐시 풀은 기동 시 고정 크기로 선할당된다")
print(BAR)
print(f"  block_size     : {cache_config.block_size} 토큰")
print(f"  num_gpu_blocks : {cache_config.num_gpu_blocks:,} 개")
print(f"  총 수용량      : {cache_config.num_gpu_blocks * cache_config.block_size:,} 토큰")
before = torch.cuda.memory_allocated()
print(f"  GPU 할당량     : {before / 1024**2:,.1f} MiB  (이후 변하지 않는다)")

# ─────────────────────────────────────────────────────────────
print("\n" + BAR)
print("[2] 요청은 max_model_len 이 아니라 '쓴 만큼'만 블록을 잡는다")
print(BAR)
print(f"  {'토큰':>5} {'블록':>5} {'ceil(n/16)':>11} {'낭비':>7}  연속할당이면")

prompts = [
    "안녕",
    "오늘 날씨는 어떤가요? 서울 기준으로 알려줘.",
    "긴 프롬프트를 만들기 위해 같은 문장을 반복한다. " * 6,
]
for i, prompt in enumerate(prompts):
    tag = f"L{i}"
    n = len(tokenizer.encode(prompt))
    add(tag, prompt)
    engine.step()                       # prefill 이 돌면서 블록이 배정된다
    got = blocks_of(tag) or []
    print(f"  {n:5d} {len(got):5d} {math.ceil(n / BLK):11d} {len(got) * BLK - n:6d}칸"
          f"  {MAX_LEN // BLK}블록 고정")
    drain()

# ─────────────────────────────────────────────────────────────
print("\n" + BAR)
print("[3] prefix 가 같으면 같은 물리 블록을 공유한다")
print(BAR)

shared = "You are a helpful assistant. " * 10      # 두 요청이 공유할 앞부분

add("PA", shared + "질문 A 입니다.")
engine.step()
blocks_a = blocks_of("PA")
drain()

add("PB", shared + "완전히 다른 질문 B 입니다.")
engine.step()
blocks_b = blocks_of("PB")
drain()

print(f"  공유 prefix : {len(tokenizer.encode(shared))} 토큰")
print(f"  A 블록 : {blocks_a}")
print(f"  B 블록 : {blocks_b}")
if blocks_a and blocks_b:
    same = [x for x, y in zip(blocks_a, blocks_b) if x == y]
    print(f"  앞에서부터 일치 : {same}"
          f"  ({len(same)}개 = {len(same) * BLK} 토큰 재계산 생략)")

# ─────────────────────────────────────────────────────────────
print("\n" + BAR)
print("[4] 그동안 GPU 메모리 총량은 변하지 않았다")
print(BAR)
after = torch.cuda.memory_allocated()
print(f"  검증 전 {before / 1024**2:,.1f} MiB"
      f" -> 검증 후 {after / 1024**2:,.1f} MiB"
      f"  (차이 {(after - before) / 1024**2:+.1f} MiB)")
print("  -> 풀은 고정, 그 안에서 블록 소유권만 오간다")

"""PagedAttention 검증 2 — 흩어진 빈칸을 주워 담는가

일부러 '구멍'을 만든다.
요청 셋을 나란히 띄우고, 가운데 것만 먼저 끝내면 그 자리가 빈다.
그 뒤 구멍보다 큰 요청을 넣어 어디서 블록을 받아 오는지 본다.
"""
import os

os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
os.environ["VLLM_USE_FLASHINFER_SAMPLER"] = "0"
os.environ["TORCHDYNAMO_DISABLE"] = "1"

from vllm import LLM, SamplingParams

BLK = 16
POOL = 80            # 풀을 좁혀야 구멍을 재사용할 수밖에 없다
BAR = "=" * 74

llm = LLM(model="HuggingFaceTB/SmolLM-135M",
          max_model_len=512,
          block_size=BLK,
          num_gpu_blocks_override=POOL,     # 기본 4,361개는 너무 넉넉하다
          enable_prefix_caching=False,      # 공유를 꺼야 파편화만 보인다
          gpu_memory_utilization=0.35,
          enforce_eager=True)

engine = llm.llm_engine
scheduler = engine.engine_core.engine_core.scheduler
manager = scheduler.kv_cache_manager.coordinator.single_type_managers[0]


def blocks_of(tag):
    for request_id, blocks in manager.req_to_blocks.items():
        if request_id == tag or request_id.startswith(tag + "-"):
            return [b.block_id for b in blocks]
    return None


def add(tag, prompt, max_tokens):
    engine.add_request(tag, prompt,
                       SamplingParams(temperature=0, max_tokens=max_tokens))


def contiguous(xs):
    s = sorted(xs)
    return s == list(range(s[0], s[-1] + 1))


print("\n" + BAR)
print(f"[1] 요청 셋을 띄운다 — 풀 {POOL}블록, prefix 공유 끔")
print(BAR)

body = "요청 %s 의 본문입니다. 블록을 여러 개 차지하도록 길게 씁니다. "
add("R0", (body % "A") * 5, max_tokens=400)   # 오래 남는다
add("R1", (body % "B") * 5, max_tokens=2)     # 먼저 끝난다 -> 구멍이 된다
add("R2", (body % "C") * 5, max_tokens=400)   # 오래 남는다
for _ in range(2):
    engine.step()

occupied = {i: blocks_of(f"R{i}") for i in range(3)}
for i in range(3):
    print(f"  R{i} : {occupied[i]}")

hole = list(occupied[1])
used = sum(len(v) for v in occupied.values() if v)
print(f"\n  사용 {used} / 전체 {POOL}  ->  아직 안 쓴 블록 {POOL - used}개")

print("\n" + BAR)
print("[2] 가운데 R1 이 끝나면서 구멍이 생긴다")
print(BAR)
for _ in range(20):
    engine.step()
    if blocks_of("R1") is None:          # 블록 테이블에서 사라지면 반환된 것
        break
print(f"  구멍 {len(hole)}블록 반환 : {hole}")
print(f"  가용 = 구멍 {len(hole)} + 미사용 {POOL - used} = {len(hole) + POOL - used}블록")
print(f"  다만 가장 긴 '연속' 구간은 {max(len(hole), POOL - used)}블록뿐이다")

print("\n" + BAR)
print("[3] 구멍보다 큰 요청을 넣는다")
print(BAR)
need = "새 요청. 흩어진 빈칸을 모두 써야 들어갈 만큼 길게 만든 프롬프트입니다. " * 5
add("NEW", need, max_tokens=4)

got = []
for _ in range(15):
    engine.step()
    got = blocks_of("NEW") or []
    if got:
        break

print(f"  NEW 가 받은 블록 ({len(got)}개) :")
print(f"    {got}")
if got:
    reused = sorted(set(got) & set(hole))
    fresh = sorted(set(got) - set(hole))
    print(f"\n  구멍에서 재사용 : {len(reused)}개 {reused}")
    print(f"  미사용분에서    : {len(fresh)}개 {fresh}")
    print(f"  연속인가? {'예' if contiguous(got) else '아니오 — 흩어진 블록을 이어 붙였다'}")
    if reused and fresh:
        print("\n  -> 연속 할당이었다면 이만한 연속 공간이 없어 거절됐을 것이다")