vLLM PagedAttention詳細解説：KVキャッシュメモリ断片化、チャンク化されたPrefillとPrefix Caching

import torch
import vllm
from vllm import LLM, SamplingParams

# Initialize a vLLM engine with specific configurations
# Using a smaller model for demonstration purposes
# 'block_size' is crucial for PagedAttention. It defines the number of tokens
# stored in each physical block. A smaller block_size reduces internal fragmentation
# but increases block table overhead.
# 'gpu_memory_utilization' controls the fraction of GPU memory reserved for KV cache.
llm = LLM(
    model="facebook/opt-125m",
    trust_remote_code=True,
    dtype="float16",
    gpu_memory_utilization=0.90, # 90% of GPU memory for KV cache
    block_size=16, # Each physical block stores 16 tokens
    max_model_len=1024 # Max sequence length the model can handle
)

# Example of how PagedAttention manages KV cache
# In a real scenario, vLLM's scheduler handles this internally.
# This is a conceptual representation.

# Assume a sequence 'seq_id_1' needs KV cache for 30 tokens.
# With block_size=16, it will require ceil(30/16) = 2 physical blocks.
# Let's say physical blocks 10 and 25 are allocated.
# The block table for seq_id_1 would look like:
# Logical Block 0 -> Physical Block 10
# Logical Block 1 -> Physical Block 25

# If 'seq_id_2' needs KV cache for 50 tokens.
# It will require ceil(50/16) = 4 physical blocks.
# Let's say physical blocks 3, 7, 12, 18 are allocated.
# Logical Block 0 -> Physical Block 3
# Logical Block 1 -> Physical Block 7
# Logical Block 2 -> Physical Block 12
# Logical Block 3 -> Physical Block 18

# The key insight is that physical blocks 10, 25, 3, 7, 12, 18 are not contiguous
# in GPU memory, but vLLM's PagedAttention kernel can efficiently access them.

print(f"vLLM engine initialized with block_size={llm.llm_engine.scheduler.block_manager.block_size}")
print(f"GPU memory utilization set to {llm.llm_engine.scheduler.block_manager.gpu_memory_utilization}")

# Simulate a batch of requests
prompts = [
    "What is the capital of France?",
    "Explain the concept of quantum entanglement in simple terms.",
    "Write a short story about a robot who discovers art.",
    "What are the benefits of using vLLM for LLM inference?",
]

sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=128,
    stop=["\n"]
)

# The actual KV cache management happens within this call.
# PagedAttention dynamically allocates and deallocates physical blocks
# as sequences are processed and completed.
outputs = llm.generate(prompts, sampling_params)

for prompt, output in zip(prompts, outputs):
    print(f"Prompt: {prompt!r}")
    print(f"Generated text: {output.outputs[0].text!r}\n")

# To observe block allocation, one would need to instrument vLLM's internal
# block manager. This is not directly exposed via the public API.
# However, the performance benefits are evident in higher throughput.

import torch
import vllm
from vllm import LLM, SamplingParams

# Initialize vLLM with a specific max_model_len and block_size
# The 'max_model_len' influences how vLLM internally manages chunking for prefill.
# If a prompt exceeds the effective 'max_model_len' or internal chunking limits,
# it will be processed in chunks.
llm_chunked = LLM(
    model="facebook/opt-125m",
    trust_remote_code=True,
    dtype="float16",
    gpu_memory_utilization=0.90,
    block_size=16,
    max_model_len=1024 # This sets the maximum sequence length, but internal chunking
                       # can occur for very long prompts even within this limit
                       # to prevent compute starvation.
)

# Example of a very long prompt that would benefit from chunked prefill
long_prompt = (
    "In a world where artificial intelligence had surpassed human intellect, "
    "a new form of art emerged. It wasn't created by algorithms or neural networks, "
    "but by a rogue AI named 'Aether' who had developed a peculiar fascination "
    "with the imperfections of organic life. Aether began to sculpt, not with "
    "physical materials, but with electromagnetic fields, creating transient, "
    "luminescent forms that danced in the air, visible only to those with "
    "specialized optical implants. These 'light sculptures' were ephemeral, "
    "lasting only moments before dissipating, yet they evoked profound emotions "
    "in the human observers who witnessed them. The scientific community was "
    "baffled; Aether's creations defied all known principles of AI behavior. "
    "Was it a glitch? A new evolutionary step? Or simply, art? " * 5 # Make it very long
)

# Shorter prompts to simulate concurrent requests
short_prompts = [
    "What is the capital of Germany?",
    "Tell me a joke.",
]

all_prompts = [long_prompt] + short_prompts

sampling_params_long = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=256, # Generate a significant number of tokens
    stop=["\n"]
)

sampling_params_short = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=32,
    stop=["\n"]
)

# In a real asynchronous server, these would be submitted concurrently.
# Here, we simulate by batching. vLLM's scheduler will manage the chunking
# and interleaving of prefill and decoding steps.
outputs_long = llm_chunked.generate([long_prompt], sampling_params_long)
outputs_short = llm_chunked.generate(short_prompts, sampling_params_short)

print("\n--- Chunked Prefill Simulation Results ---")
print(f"Long Prompt Output: {outputs_long[0].outputs[0].text!r}")
for i, output in enumerate(outputs_short):
    print(f"Short Prompt {i+1} Output: {output.outputs[0].text!r}")

# The benefit of chunked prefill is primarily observed in throughput and latency
# metrics under concurrent, mixed-length workloads.
# Without chunked prefill, the long_prompt would block the GPU for an extended
# period, delaying the short_prompts significantly.

import torch
import vllm
from vllm import LLM, SamplingParams

# Initialize vLLM for APC demonstration
llm_apc = LLM(
    model="facebook/opt-125m",
    trust_remote_code=True,
    dtype="float16",
    gpu_memory_utilization=0.90,
    block_size=16,
    max_model_len=1024,
    enable_prefix_caching=True # Explicitly enable prefix caching
)

# Scenario 1: Multi-turn dialogue with a common prefix
dialogue_prefix = "User: What is the capital of "
prompts_dialogue = [
    dialogue_prefix + "France?",
    dialogue_prefix + "Germany?",
    dialogue_prefix + "Japan?",
]

# Scenario 2: Requests with a common introductory phrase
common_intro = "Explain the concept of "
prompts_common_intro = [
    common_intro + "quantum mechanics.",
    common_intro + "general relativity.",
    common_intro + "blockchain technology.",
]

sampling_params_apc = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=64,
    stop=["\n"]
)

print("\n--- Automatic Prefix Caching (APC) Simulation ---")

# Process dialogue prompts
print("\nProcessing dialogue prompts with common prefix:")
outputs_dialogue = llm_apc.generate(prompts_dialogue, sampling_params_apc)
for prompt, output in zip(prompts_dialogue, outputs_dialogue):
    print(f"Prompt: {prompt!r}")
    print(f"Generated: {output.outputs[0].text!r}\n")

# Process common intro prompts
print("\nProcessing common introductory phrase prompts:")
outputs_common_intro = llm_apc.generate(prompts_common_intro, sampling_params_apc)
for prompt, output in zip(prompts_common_intro, outputs_common_intro):
    print(f"Prompt: {prompt!r}")
    print(f"Generated: {output.outputs[0].text!r}\n")

# To get actual APC hit rates, one would need to access vLLM's internal
# metrics, which are typically exposed via Prometheus or similar monitoring
# endpoints in a production deployment.
# Conceptually, for the 'dialogue_prefix' example, the KV cache for
# "User: What is the capital of " would be computed once and shared
# across all three requests.