51. vLLM 核心模块逐文件：engine.py

class LLMEngine:
    """The main engine for vLLM."""
    
    def __init__(self, model_config: ModelConfig, cache_config: CacheConfig,
                 parallel_config: ParallelConfig, scheduler_config: SchedulerConfig,
                 device_config: DeviceConfig, lora_config: Optional[LoRAConfig] = None,
                 distributed_init_method: Optional[str] = None):
        """Initialize the LLMEngine.
        
        Args:
            model_config: The configuration for the model.
            cache_config: The configuration for the KV cache.
            parallel_config: The configuration for parallel execution.
            scheduler_config: The configuration for the scheduler.
            device_config: The configuration for the device.
            lora_config: The configuration for LoRA adapters.
            distributed_init_method: The method for initializing distributed training.
        """
        # 初始化分布式环境
        if parallel_config.world_size > 1:
            self._init_distributed(distributed_init_method)
        
        # 创建设备上下文
        self.device_config = device_config
        self.parallel_config = parallel_config
        
        # 初始化模型运行器
        self.model_runner = ModelRunner(
            model_config=model_config,
            parallel_config=parallel_config,
            device_config=device_config,
            lora_config=lora_config,
        )
        
        # 初始化块管理器
        self.block_manager = BlockManager(
            cache_config=cache_config,
            block_size=self.model_runner.get_block_size(),
            num_gpus=parallel_config.world_size,
            device_config=device_config,
        )
        
        # 初始化调度器
        self.scheduler = Scheduler(
            scheduler_config=scheduler_config,
            block_manager=self.block_manager,
            num_gpus=parallel_config.world_size,
        )
        
        # 初始化请求ID计数器
        self.request_id_counter = 0
        
        # 初始化统计信息
        self.stats = EngineStats()

async def add_request(self, request_id: Optional[int] = None, prompt: Optional[str] = None,
                      prompt_token_ids: Optional[List[int]] = None, sampling_params: SamplingParams = None,
                      lora_request: Optional[LoRARequest] = None, arrival_time: Optional[float] = None,
                      **kwargs) -> int:
    """Add a new request to the engine.
    
    Args:
        request_id: The ID of the request. If None, a new ID will be generated.
        prompt: The prompt string. Either prompt or prompt_token_ids must be provided.
        prompt_token_ids: The token IDs of the prompt. Either prompt or prompt_token_ids must be provided.
        sampling_params: The sampling parameters for text generation.
        lora_request: The LoRA request for this generation.
        arrival_time: The arrival time of the request. If None, current time will be used.
        **kwargs: Additional arguments for the request.
    
    Returns:
        The ID of the added request.
    """
    # 生成请求ID
    if request_id is None:
        request_id = self.request_id_counter
        self.request_id_counter += 1
    
    # 计算请求到达时间
    if arrival_time is None:
        arrival_time = time.time()
    
    # 处理提示文本，转换为token IDs
    if prompt_token_ids is None:
        if prompt is None:
            raise ValueError("Either prompt or prompt_token_ids must be provided.")
        prompt_token_ids = self.model_runner.tokenize_prompt(prompt)
    
    # 创建请求对象
    request = Request(
        request_id=request_id,
        prompt=prompt,
        prompt_token_ids=prompt_token_ids,
        sampling_params=sampling_params,
        lora_request=lora_request,
        arrival_time=arrival_time,
        **kwargs
    )
    
    # 将请求添加到调度器
    await self.scheduler.add_request(request)
    
    # 更新统计信息
    self.stats.num_requests += 1
    
    return request_id

async def step(self) -> List[RequestOutput]:
    """Run one step of the engine.
    
    Returns:
        A list of completed requests.
    """
    # 从调度器获取当前批次
    batch = await self.scheduler.schedule()
    if batch is None:
        return []
    
    # 执行模型前向计算
    outputs = await self.model_runner.execute_model(batch)
    
    # 更新块管理器
    self.block_manager.update_batches(batch, outputs)
    
    # 处理采样结果
    completed_requests = await self.scheduler.process_model_outputs(batch, outputs)
    
    # 更新统计信息
    self.stats.update_step_stats(batch, outputs)
    
    return completed_requests

def _init_distributed(self, distributed_init_method: Optional[str] = None):
    """Initialize distributed training."""
    if self.parallel_config.tensor_parallel_size > 1:
        # 初始化张量并行组
        self.tp_rank = dist.get_rank(group=dist.new_group(ranks=list(range(self.parallel_config.tensor_parallel_size))))
        self.tp_world_size = self.parallel_config.tensor_parallel_size
    else:
        self.tp_rank = 0
        self.tp_world_size = 1
    
    if self.parallel_config.pipeline_parallel_size > 1:
        # 初始化流水线并行组
        self.pp_rank = dist.get_rank(group=dist.new_group(ranks=list(range(self.parallel_config.pipeline_parallel_size))))
        self.pp_world_size = self.parallel_config.pipeline_parallel_size
    else:
        self.pp_rank = 0
        self.pp_world_size = 1

async def _distributed_schedule(self, batch: Batch) -> Batch:
    """Schedule the batch for distributed execution."""
    if self.tp_world_size > 1:
        # 张量并行场景下的批处理分割
        batch = self._split_batch_for_tensor_parallel(batch)
    
    if self.pp_world_size > 1:
        # 流水线并行场景下的批处理调度
        batch = await self._schedule_for_pipeline_parallel(batch)
    
    return batch

def _adjust_batch_size(self, batch: Batch) -> Batch:
    """Adjust the batch size based on current GPU memory usage."""
    current_memory = self._get_current_gpu_memory()
    max_memory = self._get_max_gpu_memory()
    
    # 根据内存使用情况调整批次大小
    if current_memory > max_memory * 0.8:
        # 内存紧张，减小批次大小
        batch = self._reduce_batch_size(batch, target_ratio=0.5)
    elif current_memory < max_memory * 0.5:
        # 内存充足，尝试增大批次大小
        batch = await self._increase_batch_size(batch, target_ratio=1.5)
    
    return batch

def _optimize_kv_cache(self, batch: Batch) -> None:
    """Optimize KV cache usage for the current batch."""
    # 合并相邻的KV缓存块
    self.block_manager.merge_adjacent_blocks()
    
    # 释放不再使用的KV缓存块
    self.block_manager.free_unused_blocks()
    
    # 压缩KV缓存（如果启用）
    if self.config.enable_kv_cache_compression:
        self.block_manager.compress_kv_cache()

from vllm.engine import LLMEngine
from vllm.config import ModelConfig, CacheConfig, ParallelConfig, SchedulerConfig, DeviceConfig
from vllm.sampling_params import SamplingParams

# 配置模型参数
model_config = ModelConfig(
    model="meta-llama/Llama-2-70b-hf",
    dtype="float16",
    trust_remote_code=True,
)

# 配置缓存参数
cache_config = CacheConfig(
    block_size=16,
    gpu_memory_utilization=0.9,
    swap_space=4,
)

# 配置并行参数
parallel_config = ParallelConfig(
    tensor_parallel_size=8,
    pipeline_parallel_size=1,
)

# 配置调度器参数
scheduler_config = SchedulerConfig(
    max_num_seqs=256,
    max_model_len=4096,
)

# 配置设备参数
device_config = DeviceConfig(
    device="cuda",
    seed=42,
)

# 创建LLMEngine实例
engine = LLMEngine(
    model_config=model_config,
    cache_config=cache_config,
    parallel_config=parallel_config,
    scheduler_config=scheduler_config,
    device_config=device_config,
)

# 定义采样参数
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=128,
)

# 添加请求
request_id = await engine.add_request(
    prompt="Hello, my name is",
    sampling_params=sampling_params,
)

# 执行引擎循环
completed_requests = []
while len(completed_requests) == 0:
    completed_requests = await engine.step()

# 输出结果
for request in completed_requests:
    print(f"Request {request.request_id} completed:")
    print(f"Generated text: {request.outputs[0].text}")
    print(f"Tokens generated: {len(request.outputs[0].token_ids)}")

Request 0 completed:
Generated text: Hello, my name is John. I'm a software engineer with over 10 years of experience in building scalable web applications. I specialize in Python, JavaScript, and cloud computing. In my free time, I enjoy hiking, reading, and spending time with my family.
Tokens generated: 56

import os
import torch
from vllm.engine import LLMEngine
from vllm.config import ModelConfig, CacheConfig, ParallelConfig, SchedulerConfig, DeviceConfig

# 设置分布式环境变量
os.environ["RANK"] = str(torch.distributed.get_rank())
os.environ["WORLD_SIZE"] = str(torch.distributed.get_world_size())
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "29500"

# 初始化分布式环境
torch.distributed.init_process_group(backend="nccl")

# 配置模型参数
model_config = ModelConfig(
    model="meta-llama/Llama-2-70b-hf",
    dtype="float16",
    trust_remote_code=True,
)

# 配置缓存参数
cache_config = CacheConfig(
    block_size=16,
    gpu_memory_utilization=0.9,
)

# 配置并行参数（8卡张量并行）
parallel_config = ParallelConfig(
    tensor_parallel_size=8,
    pipeline_parallel_size=1,
)

# 配置调度器参数
scheduler_config = SchedulerConfig(
    max_num_seqs=256,
    max_model_len=4096,
)

# 配置设备参数
device_config = DeviceConfig(
    device="cuda",
    seed=42,
)

# 创建分布式LLMEngine实例
engine = LLMEngine(
    model_config=model_config,
    cache_config=cache_config,
    parallel_config=parallel_config,
    scheduler_config=scheduler_config,
    device_config=device_config,
    distributed_init_method="env://",
)

# 主进程处理请求和输出
if torch.distributed.get_rank() == 0:
    from vllm.sampling_params import SamplingParams
    
    # 定义采样参数
sampling_params = SamplingParams(
        temperature=0.7,
        top_p=0.95,
        max_tokens=128,
    )
    
    # 添加多个请求
    request_ids = []
    prompts = [
        "Hello, how are you?",
        "What's the capital of France?",
        "Explain quantum computing in simple terms.",
    ]
    
    for prompt in prompts:
        request_id = await engine.add_request(
            prompt=prompt,
            sampling_params=sampling_params,
        )
        request_ids.append(request_id)
    
    # 执行引擎循环，直到所有请求完成
    completed_requests = []
    while len(completed_requests) < len(request_ids):
        new_completed = await engine.step()
        completed_requests.extend(new_completed)
    
    # 输出结果
    for request in completed_requests:
        print(f"Request {request.request_id} completed:")
        print(f"Prompt: {request.prompt}")
        print(f"Generated text: {request.outputs[0].text}")
        print(f"Tokens generated: {len(request.outputs[0].token_ids)}")
        print()

# 销毁分布式进程组
torch.distributed.destroy_process_group()

python -m torch.distributed.run --nproc_per_node=8 distributed_engine_example.py

Request 0 completed:
Prompt: Hello, how are you?
Generated text: Hello, I'm doing well, thank you! How can I assist you today?
Tokens generated: 18

Request 1 completed:
Prompt: What's the capital of France?
Generated text: The capital of France is Paris.
Tokens generated: 8

Request 2 completed:
Prompt: Explain quantum computing in simple terms.
Generated text: Quantum computing is a type of computing that uses the principles of quantum mechanics to process information. Unlike classical computers, which use bits (0s and 1s) to store and process data, quantum computers use quantum bits, or qubits, which can exist in multiple states at once. This allows quantum computers to perform certain calculations much faster than classical computers.
Tokens generated: 57

# 安装vLLM
pip install vllm

# 安装分布式依赖
pip install ray torch torchvision torchaudio

# 安装监控工具
pip install prometheus-client