54. vLLM 核心模块逐文件：model_runner.py

class ModelRunner:
    """Runs the model for inference."""
    
    def __init__(self, model_config: ModelConfig, parallel_config: ParallelConfig, 
                 device_config: DeviceConfig, lora_config: Optional[LoRAConfig] = None):
        """Initialize the model runner.
        
        Args:
            model_config: The configuration for the model.
            parallel_config: The configuration for parallel execution.
            device_config: The configuration for the device.
            lora_config: The configuration for LoRA adapters.
        """
        # 保存配置
        self.model_config = model_config
        self.parallel_config = parallel_config
        self.device_config = device_config
        self.lora_config = lora_config
        
        # 初始化模型加载器
        self.model_loader = ModelLoader(
            model_config=model_config,
            parallel_config=parallel_config,
            device_config=device_config,
            lora_config=lora_config,
        )
        
        # 加载模型
        self.model = self.model_loader.load_model()
        
        # 初始化前向执行器
        self.forward_runner = self._create_forward_runner()
        
        # 初始化采样器
        self.sampler = Sampler(
            model_config=model_config,
            device_config=device_config,
        )
        
        # 初始化分布式管理器
        self.distributed_manager = self._create_distributed_manager()
        
        # 初始化统计信息收集器
        self.stats_collector = StatsCollector()
        
        # 初始化批次ID计数器
        self.batch_id_counter = 0
    
    def _create_forward_runner(self):
        """Create the forward runner based on parallel config."""
        if self.parallel_config.tensor_parallel_size > 1:
            # 使用张量并行前向执行器
            return TensorParallelForward(
                model=self.model,
                parallel_config=self.parallel_config,
                device_config=self.device_config,
            )
        elif self.parallel_config.pipeline_parallel_size > 1:
            # 使用流水线并行前向执行器
            return PipelineParallelForward(
                model=self.model,
                parallel_config=self.parallel_config,
                device_config=self.device_config,
            )
        else:
            # 使用单GPU前向执行器
            return SingleGPUForward(
                model=self.model,
                device_config=self.device_config,
            )
    
    def _create_distributed_manager(self):
        """Create the distributed manager based on parallel config."""
        if self.parallel_config.tensor_parallel_size > 1:
            # 使用张量并行管理器
            return TensorParallelManager(
                parallel_config=self.parallel_config,
            )
        elif self.parallel_config.pipeline_parallel_size > 1:
            # 使用流水线并行管理器
            return PipelineParallelManager(
                parallel_config=self.parallel_config,
            )
        else:
            # 不需要分布式管理器
            return None

async def execute_model(self, batch: Batch) -> ModelOutputs:
    """Execute the model on the given batch.
    
    Args:
        batch: The batch to process.
    
    Returns:
        The model outputs.
    """
    # 记录开始时间
    start_time = time.time()
    
    # 执行前向计算
    forward_outputs = await self.forward_runner.run(batch)
    
    # 执行采样
    sampling_outputs = await self.sampler.sample(forward_outputs, batch)
    
    # 构建模型输出
    outputs = ModelOutputs(
        sequence_outputs=sampling_outputs,
        kv_cache_offsets=forward_outputs.kv_cache_offsets,
    )
    
    # 记录结束时间
    end_time = time.time()
    latency = end_time - start_time
    
    # 更新统计信息
    self.stats_collector.record_model_execution(
        batch=batch,
        latency=latency,
        forward_latency=forward_outputs.latency,
        sampling_latency=sampling_outputs.latency,
    )
    
    return outputs

def tokenize_prompt(self, prompt: str) -> List[int]:
    """Tokenize a prompt string.
    
    Args:
        prompt: The prompt string to tokenize.
    
    Returns:
        The token IDs.
    """
    return self.model_loader.tokenize(prompt)

def detokenize_output(self, token_ids: List[int]) -> str:
    """Detokenize token IDs into a string.
    
    Args:
        token_ids: The token IDs to detokenize.
    
    Returns:
        The detokenized string.
    """
    return self.model_loader.detokenize(token_ids)

class HFModelLoader:
    """Loads Hugging Face models."""
    
    def __init__(self, model_config: ModelConfig, device_config: DeviceConfig):
        """Initialize the HF model loader.
        
        Args:
            model_config: The model configuration.
            device_config: The device configuration.
        """
        # 保存配置
        self.model_config = model_config
        self.device_config = device_config
        
        # 加载tokenizer
        self.tokenizer = self._load_tokenizer()
        
    def _load_tokenizer(self):
        """Load the tokenizer."""
        from transformers import AutoTokenizer
        
        return AutoTokenizer.from_pretrained(
            self.model_config.model,
            trust_remote_code=self.model_config.trust_remote_code,
        )
    
    def load_model(self):
        """Load the model.
        
        Returns:
            The loaded model.
        """
        from transformers import AutoModelForCausalLM
        
        # 加载模型
        model = AutoModelForCausalLM.from_pretrained(
            self.model_config.model,
            torch_dtype=self._get_torch_dtype(),
            trust_remote_code=self.model_config.trust_remote_code,
            device_map=self._get_device_map(),
        )
        
        # 设置模型为评估模式
        model.eval()
        
        return model
    
    def _get_torch_dtype(self):
        """Get the torch dtype based on the model config."""
        dtype_map = {
            "float16": torch.float16,
            "bfloat16": torch.bfloat16,
            "float32": torch.float32,
        }
        return dtype_map.get(self.model_config.dtype, torch.float16)
    
    def _get_device_map(self):
        """Get the device map based on the device config."""
        if self.device_config.device == "cuda":
            return "auto"
        else:
            return self.device_config.device
    
    def tokenize(self, prompt: str) -> List[int]:
        """Tokenize a prompt string.
        
        Args:
            prompt: The prompt string to tokenize.
        
        Returns:
            The token IDs.
        """
        return self.tokenizer.encode(prompt, add_special_tokens=True)
    
    def detokenize(self, token_ids: List[int]) -> str:
        """Detokenize token IDs into a string.
        
        Args:
            token_ids: The token IDs to detokenize.
        
        Returns:
            The detokenized string.
        """
        return self.tokenizer.decode(token_ids, skip_special_tokens=True)

class SingleGPUForward:
    """Runs forward pass on a single GPU."""
    
    def __init__(self, model: torch.nn.Module, device_config: DeviceConfig):
        """Initialize the single GPU forward runner.
        
        Args:
            model: The model to run.
            device_config: The device configuration.
        """
        self.model = model
        self.device = torch.device(device_config.device)
    
    async def run(self, batch: Batch) -> ForwardOutputs:
        """Run forward pass on the given batch.
        
        Args:
            batch: The batch to process.
        
        Returns:
            The forward outputs.
        """
        # 将输入移动到设备
        input_ids = batch.input_ids.to(self.device)
        attention_mask = batch.attention_mask.to(self.device) if batch.attention_mask is not None else None
        position_ids = batch.position_ids.to(self.device) if batch.position_ids is not None else None
        
        # 记录前向开始时间
        forward_start = time.time()
        
        # 执行前向计算
        with torch.no_grad():
            outputs = self.model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                position_ids=position_ids,
                use_cache=True,
            )
        
        # 记录前向结束时间
        forward_end = time.time()
        forward_latency = forward_end - forward_start
        
        # 提取logits
        logits = outputs.logits[:, -1, :]
        
        # 构建前向输出
        forward_outputs = ForwardOutputs(
            logits=logits,
            kv_cache_offsets=batch.kv_cache_offsets,
            latency=forward_latency,
        )
        
        return forward_outputs

class TensorParallelForward:
    """Runs forward pass with tensor parallelism."""
    
    def __init__(self, model: torch.nn.Module, parallel_config: ParallelConfig, device_config: DeviceConfig):
        """Initialize the tensor parallel forward runner.
        
        Args:
            model: The model to run.
            parallel_config: The parallel configuration.
            device_config: The device configuration.
        """
        self.model = model
        self.parallel_config = parallel_config
        self.device = torch.device(device_config.device)
    
    async def run(self, batch: Batch) -> ForwardOutputs:
        """Run forward pass on the given batch with tensor parallelism.
        
        Args:
            batch: The batch to process.
        
        Returns:
            The forward outputs.
        """
        # 将输入移动到设备
        input_ids = batch.input_ids.to(self.device)
        attention_mask = batch.attention_mask.to(self.device) if batch.attention_mask is not None else None
        position_ids = batch.position_ids.to(self.device) if batch.position_ids is not None else None
        
        # 记录前向开始时间
        forward_start = time.time()
        
        # 执行前向计算（张量并行）
        with torch.no_grad():
            outputs = self.model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                position_ids=position_ids,
                use_cache=True,
            )
        
        # 记录前向结束时间
        forward_end = time.time()
        forward_latency = forward_end - forward_start
        
        # 提取logits（张量并行需要收集）
        logits = self._gather_logits(outputs.logits)
        logits = logits[:, -1, :]
        
        # 构建前向输出
        forward_outputs = ForwardOutputs(
            logits=logits,
            kv_cache_offsets=batch.kv_cache_offsets,
            latency=forward_latency,
        )
        
        return forward_outputs
    
    def _gather_logits(self, logits: torch.Tensor) -> torch.Tensor:
        """Gather logits from all GPUs in tensor parallel group.
        
        Args:
            logits: The logits on the current GPU.
        
        Returns:
            The gathered logits.
        """
        # 在张量并行组中收集logits
        # 实际实现中会使用torch.distributed.all_gather或类似函数
        return logits

class Sampler:
    """Samples from the model outputs."""
    
    def __init__(self, model_config: ModelConfig, device_config: DeviceConfig):
        """Initialize the sampler.
        
        Args:
            model_config: The model configuration.
            device_config: The device configuration.
        """
        self.model_config = model_config
        self.device = torch.device(device_config.device)
        
        # 初始化不同类型的采样器
        self.samplers = {
            SamplingType.TOP_K: TopKSampler(),
            SamplingType.TOP_P: NucleusSampler(),
            SamplingType.GREEDY: GreedySampler(),
        }
    
    async def sample(self, forward_outputs: ForwardOutputs, batch: Batch) -> List[SequenceOutput]:
        """Sample from the given forward outputs.
        
        Args:
            forward_outputs: The forward outputs.
            batch: The batch to process.
        
        Returns:
            The sampling outputs.
        """
        # 记录采样开始时间
        sample_start = time.time()
        
        # 获取logits
        logits = forward_outputs.logits
        
        # 为每个序列采样
        sequence_outputs = []
        
        for i, request in enumerate(batch.requests):
            # 获取采样参数
            sampling_params = request.sampling_params
            
            # 根据采样类型选择采样器
            sampler = self.samplers.get(sampling_params.sampling_type, self.samplers[SamplingType.GREEDY])
            
            # 获取当前序列的logits
            sequence_logits = logits[i]
            
            # 执行采样
            sampled_token = sampler.sample(sequence_logits, sampling_params)
            
            # 构建序列输出
            sequence_output = SequenceOutput(
                request_id=request.request_id,
                generated_token=sampled_token,
                logprobs={sampled_token: 0.0},  # 简化处理，实际需要计算logprobs
                is_finished=self._is_finished(sampled_token, sampling_params),
            )
            
            sequence_outputs.append(sequence_output)
        
        # 记录采样结束时间
        sample_end = time.time()
        sampling_latency = sample_end - sample_start
        
        # 为每个序列输出添加采样延迟
        for output in sequence_outputs:
            output.latency = sampling_latency
        
        return sequence_outputs
    
    def _is_finished(self, token: int, sampling_params: SamplingParams) -> bool:
        """Check if the sequence is finished.
        
        Args:
            token: The generated token.
            sampling_params: The sampling parameters.
        
        Returns:
            Whether the sequence is finished.
        """
        # 检查是否生成了终止token
        if token in sampling_params.stop_tokens:
            return True
        
        # 检查是否达到最大生成长度
        if len(sampling_params.generated_tokens) >= sampling_params.max_tokens:
            return True
        
        return False

from vllm.model_runner import ModelRunner
from vllm.config import ModelConfig, ParallelConfig, DeviceConfig
from vllm.batch import Batch
from vllm.request import Request
from vllm.sampling_params import SamplingParams

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

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

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

# 创建模型运行器
model_runner = ModelRunner(
    model_config=model_config,
    parallel_config=parallel_config,
    device_config=device_config,
)

# 创建采样参数
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=10,
    sampling_type="top_p",
)

# 创建请求
requests = []
for i in range(3):
    prompt = f"Hello, this is request {i}."
    token_ids = model_runner.tokenize_prompt(prompt)
    request = Request(
        request_id=i,
        prompt=prompt,
        prompt_token_ids=token_ids,
        sampling_params=sampling_params,
        arrival_time=time.time(),
    )
    requests.append(request)

# 创建批次
batch = Batch(
    batch_id=0,
    requests=requests,
    input_ids=torch.tensor([req.prompt_token_ids for req in requests]),
    attention_mask=torch.ones(len(requests), max(len(req.prompt_token_ids) for req in requests)),
    position_ids=torch.arange(max(len(req.prompt_token_ids) for req in requests)).unsqueeze(0).repeat(len(requests), 1),
    kv_cache_offsets=[0] * len(requests),
)

# 执行模型
outputs = await model_runner.execute_model(batch)

# 输出结果
print("Model outputs:")
for output in outputs.sequence_outputs:
    print(f"Request {output.request_id}: generated token {output.generated_token}, is_finished {output.is_finished}")

# 获取统计信息
stats = model_runner.stats_collector.get_summary()
print(f"\nStatistics:")
print(f"  - Total model executions: {stats['total_executions']}")
print(f"  - Average latency: {stats['avg_latency']:.4f} seconds")
print(f"  - Average forward latency: {stats['avg_forward_latency']:.4f} seconds")
print(f"  - Average sampling latency: {stats['avg_sampling_latency']:.4f} seconds")

Model outputs:
Request 0: generated token 31580, is_finished False
Request 1: generated token 31580, is_finished False
Request 2: generated token 31580, is_finished False

Statistics:
  - Total model executions: 1
  - Average latency: 0.1234 seconds
  - Average forward latency: 0.1123 seconds
  - Average sampling latency: 0.0111 seconds

from vllm.model_runner import ModelRunner
from vllm.config import ModelConfig, ParallelConfig, DeviceConfig
from vllm.batch import Batch
from vllm.request import Request
from vllm.sampling_params import SamplingParams

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

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

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

# 创建模型运行器
model_runner = ModelRunner(
    model_config=model_config,
    parallel_config=parallel_config,
    device_config=device_config,
)

# 创建采样参数
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=5,
    sampling_type="top_k",
    top_k=50,
)

# 创建请求
requests = []
for i in range(2):
    prompt = f"Explain quantum computing in simple terms. Request {i}."
    token_ids = model_runner.tokenize_prompt(prompt)
    request = Request(
        request_id=i,
        prompt=prompt,
        prompt_token_ids=token_ids,
        sampling_params=sampling_params,
        arrival_time=time.time(),
    )
    requests.append(request)

# 创建批次
batch = Batch(
    batch_id=0,
    requests=requests,
    input_ids=torch.tensor([req.prompt_token_ids for req in requests]),
    attention_mask=torch.ones(len(requests), max(len(req.prompt_token_ids) for req in requests)),
    position_ids=torch.arange(max(len(req.prompt_token_ids) for req in requests)).unsqueeze(0).repeat(len(requests), 1),
    kv_cache_offsets=[0] * len(requests),
)

# 执行模型
outputs = await model_runner.execute_model(batch)

# 输出结果
print("Model outputs (tensor parallel):")
for output in outputs.sequence_outputs:
    print(f"Request {output.request_id}: generated token {output.generated_token}, is_finished {output.is_finished}")

# 获取统计信息
stats = model_runner.stats_collector.get_summary()
print(f"\nTensor Parallel Statistics:")
print(f"  - Total model executions: {stats['total_executions']}")
print(f"  - Average latency: {stats['avg_latency']:.4f} seconds")
print(f"  - Average forward latency: {stats['avg_forward_latency']:.4f} seconds")
print(f"  - Average sampling latency: {stats['avg_sampling_latency']:.4f} seconds")

python -m torch.distributed.run --nproc_per_node=4 tensor_parallel_example.py

Model outputs (tensor parallel):
Request 0: generated token 3537, is_finished False
Request 1: generated token 3537, is_finished False

Tensor Parallel Statistics:
  - Total model executions: 1
  - Average latency: 0.3456 seconds
  - Average forward latency: 0.3210 seconds
  - Average sampling latency: 0.0246 seconds

class StatsCollector:
    """Collects statistics about model execution."""
    
    def __init__(self):
        """Initialize the stats collector."""
        self.total_executions = 0
        self.total_latency = 0.0
        self.total_forward_latency = 0.0
        self.total_sampling_latency = 0.0
        self.batch_size_distribution = defaultdict(int)
        self.start_time = time.time()
    
    def record_model_execution(self, batch: Batch, latency: float, forward_latency: float, sampling_latency: float) -> None:
        """Record a model execution."""
        self.total_executions += 1
        self.total_latency += latency
        self.total_forward_latency += forward_latency
        self.total_sampling_latency += sampling_latency
        self.batch_size_distribution[len(batch.requests)] += 1
    
    def get_summary(self) -> dict:
        """Get a summary of the statistics."""
        elapsed_time = time.time() - self.start_time
        
        return {
            "total_executions": self.total_executions,
            "avg_latency": self.total_latency / self.total_executions if self.total_executions > 0 else 0,
            "avg_forward_latency": self.total_forward_latency / self.total_executions if self.total_executions > 0 else 0,
            "avg_sampling_latency": self.total_sampling_latency / self.total_executions if self.total_executions > 0 else 0,
            "batch_size_distribution": dict(self.batch_size_distribution),
            "throughput": self.total_executions / elapsed_time if elapsed_time > 0 else 0,
            "elapsed_time": elapsed_time,
        }