sglang框架源码笔记

整体架构

这张图片展示了 SGLang 系统的架构概览，主要由 客户端 和 服务器端 两大部分组成，展示了其各个组件之间的交互。以下是各部分的详细解释：

1. 客户端（Client）：

• Native generation API：用于本地生成模型输出的 API，允许用户直接与模型进行交互，获取生成的文本或其他输出。
• OpenAI-compatible API：兼容 OpenAI API 的接口，允许用户使用与 OpenAI 模型相同的接口调用模型，这使得与其他 OpenAI 模型的兼容性更高。
• Structured language frontend：这是一个结构化语言前端，可能用于输入特定格式的请求，支持更复杂或定制化的输入方式，如通过结构化语言定义输入数据格式。

2. 服务器端（Server）：

• API server：处理客户端请求的服务器，接收来自客户端的请求并协调后续的处理。
• Tokenizer：负责对输入数据进行标记化（tokenization），将原始文本转化为模型能够理解的 token。模型推理通常需要标记化后的数据。
• Detokenizer：与 Tokenizer 对应，负责将模型生成的 tokens 转换回原始文本（反标记化）。

3. 调度器与模型工作节点（Scheduler & Model Worker）：

• KV Cache memory pool：关键值（KV）缓存内存池，每个 token 都占用一个内存页面。缓存用于加速模型计算，存储中间结果或其他必要的数据。
• Radix tree cache：可能是用来优化数据访问速度的一种树形缓存结构，特别适合快速查找和更新。
• Finite state machine analyzer：有限状态机分析器，用于分析和控制模型推理过程中可能的状态变化，确保推理过程的正确性和效率。
• Attention backend (FlashInfer and Triton)：注意力机制的后端，使用 FlashInfer 和 Triton 进行加速计算。FlashInfer 和 Triton 可能是针对 GPU 和大规模推理的加速工具，可以显著提高推理性能。

TpModelWorker类

import logging
import threading
from typing import Optional# 导入配置文件和工具函数
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.hf_transformers_utils import get_processor, get_tokenizer
from sglang.srt.managers.io_struct import (GetWeightsByNameReqInput,InitWeightsUpdateGroupReqInput,UpdateWeightFromDiskReqInput,UpdateWeightsFromDistributedReqInput,UpdateWeightsFromTensorReqInput,
)
from sglang.srt.managers.schedule_batch import ModelWorkerBatch, global_server_args_dict
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import MultiprocessingSerializer, broadcast_pyobj, set_random_seed# 设置日志记录器
logger = logging.getLogger(__name__)class TpModelWorker:"""A tensor parallel model worker (一个张量并行模型工作者)"""def __init__(self,server_args: ServerArgs,gpu_id: int,tp_rank: int,dp_rank: Optional[int],nccl_port: int,is_draft_worker: bool = False,):# 解析传入的参数，初始化张量并行工作者self.tp_rank = tp_rank# 初始化模型配置对象，用于加载模型配置self.model_config = ModelConfig((server_args.model_pathif not is_draft_workerelse server_args.speculative_draft_model_path),trust_remote_code=server_args.trust_remote_code,revision=server_args.revision,context_length=server_args.context_length,model_override_args=server_args.json_model_override_args,is_embedding=server_args.is_embedding,dtype=server_args.dtype,quantization=server_args.quantization,)# 初始化模型执行器 (ModelRunner)，负责实际的推理过程self.model_runner = ModelRunner(model_config=self.model_config,mem_fraction_static=server_args.mem_fraction_static,gpu_id=gpu_id,tp_rank=tp_rank,tp_size=server_args.tp_size,nccl_port=nccl_port,server_args=server_args,is_draft_worker=is_draft_worker,)# 如果需要跳过tokenizer初始化，则不初始化tokenizer和processorif server_args.skip_tokenizer_init:self.tokenizer = self.processor = Noneelse:# 如果模型是多模态的，使用处理器和tokenizerif self.model_config.is_multimodal:self.processor = get_processor(server_args.tokenizer_path,tokenizer_mode=server_args.tokenizer_mode,trust_remote_code=server_args.trust_remote_code,revision=server_args.revision,)self.tokenizer = self.processor.tokenizerelse:# 否则，单独初始化tokenizerself.tokenizer = get_tokenizer(server_args.tokenizer_path,tokenizer_mode=server_args.tokenizer_mode,trust_remote_code=server_args.trust_remote_code,revision=server_args.revision,)# 获取模型运行设备（GPU）self.device = self.model_runner.device# 配置与内存相关的参数self.max_total_num_tokens = self.model_runner.max_total_num_tokensself.max_prefill_tokens = server_args.max_prefill_tokensself.max_running_requests = min((self.max_total_num_tokens // 2if server_args.max_running_requests is Noneelse server_args.max_running_requests// (server_args.dp_size if server_args.enable_dp_attention else 1)),self.model_runner.req_to_token_pool.size,)self.max_req_len = min(self.model_config.context_len - 1,self.max_total_num_tokens - 1,)self.max_req_input_len = self.max_req_len - 5# 确保内存池足够大assert (self.max_req_len > 0 and self.max_req_input_len > 0), "Memory pool size is too small"# 在张量并行工作者之间同步随机种子self.random_seed = broadcast_pyobj([server_args.random_seed],self.tp_rank,self.model_runner.tp_group.cpu_group,)[0]# 设置随机种子，确保可重复性set_random_seed(self.random_seed)def get_worker_info(self):"""返回与工作者相关的信息，如最大token数量、最大请求长度、随机种子等"""return (self.max_total_num_tokens,self.max_prefill_tokens,self.max_running_requests,self.max_req_len,self.max_req_input_len,self.random_seed,self.device,global_server_args_dict,self.model_runner.req_to_token_pool.size,self.model_runner.req_to_token_pool.max_context_len,self.model_runner.token_to_kv_pool.size,)def get_pad_input_ids_func(self):"""返回用于填充输入ID的函数（如果模型有提供的话）"""return getattr(self.model_runner.model, "pad_input_ids", None)def get_tp_cpu_group(self):"""返回张量并行的CPU组，用于协调并行工作"""return self.model_runner.tp_group.cpu_groupdef get_attention_tp_cpu_group(self):"""返回注意力机制相关的张量并行CPU组"""return self.model_runner.attention_tp_group.cpu_groupdef get_memory_pool(self):"""返回工作者的内存池，用于存储tokens和key-value对"""return (self.model_runner.req_to_token_pool,self.model_runner.token_to_kv_pool,)def forward_batch_generation(self,model_worker_batch: ModelWorkerBatch,launch_done: Optional[threading.Event] = None,skip_sample: bool = False,):"""执行一批生成任务的前向传播，并根据需要采样token"""# 初始化前向批次forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)# 执行前向传播logits_output = self.model_runner.forward(forward_batch)# 如果提供了launch_done事件，标记为完成if launch_done:launch_done.set()# 如果skip_sample为True，则跳过token采样if skip_sample:next_token_ids = Noneelse:next_token_ids = self.model_runner.sample(logits_output, model_worker_batch)return logits_output, next_token_idsdef forward_batch_embedding(self, model_worker_batch: ModelWorkerBatch):"""执行一批embedding任务的前向传播，返回嵌入向量"""forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)logits_output = self.model_runner.forward(forward_batch)embeddings = logits_output.embeddingsreturn embeddingsdef update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput):"""从磁盘更新权重"""success, message = self.model_runner.update_weights_from_disk(recv_req.model_path, recv_req.load_format)return success, messagedef init_weights_update_group(self, recv_req: InitWeightsUpdateGroupReqInput):"""初始化权重更新的分布式组"""success, message = self.model_runner.init_weights_update_group(recv_req.master_address,recv_req.master_port,recv_req.rank_offset,recv_req.world_size,recv_req.group_name,recv_req.backend,)return success, messagedef update_weights_from_distributed(self, recv_req: UpdateWeightsFromDistributedReqInput):"""从分布式源更新权重"""success, message = self.model_runner.update_weights_from_distributed(recv_req.name, recv_req.dtype, recv_req.shape)return success, messagedef update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):"""从张量更新权重"""success, message = self.model_runner.update_weights_from_tensor(MultiprocessingSerializer.deserialize(recv_req.serialized_named_tensors))return success, messagedef get_weights_by_name(self, recv_req: GetWeightsByNameReqInput):"""根据权重名称获取对应的权重参数"""parameter = self.model_runner.get_weights_by_name(recv_req.name, recv_req.truncate_size)return parameter

ModelRunner类

import gc
import json
import logging
import time
from typing import List, Optional, Tupleimport torch
import torch.distributed as distfrom sglang.srt.configs.device_config import DeviceConfig
from sglang.srt.configs.load_config import LoadConfig
from sglang.srt.configs.model_config import AttentionArch, ModelConfig
from sglang.srt.distributed import (get_tp_group,init_distributed_environment,initialize_model_parallel,set_custom_all_reduce,
)
from sglang.srt.distributed.parallel_state import monkey_patch_vllm_parallel_state
from sglang.srt.layers.attention.double_sparsity_backend import DoubleSparseAttnBackend
from sglang.srt.layers.attention.flashinfer_backend import FlashInferAttnBackend
from sglang.srt.layers.attention.flashinfer_mla_backend import FlashInferMLAAttnBackend
from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBackend
from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
from sglang.srt.layers.dp_attention import (get_attention_tp_group,get_attention_tp_size,initialize_dp_attention,
)
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.sampler import Sampler
from sglang.srt.layers.torchao_utils import apply_torchao_config_to_model
from sglang.srt.lora.lora_manager import LoRAManager
from sglang.srt.managers.schedule_batch import global_server_args_dict
from sglang.srt.mem_cache.memory_pool import (DoubleSparseTokenToKVPool,MHATokenToKVPool,MLATokenToKVPool,ReqToTokenPool,
)
from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_loader import get_model
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.torch_memory_saver_adapter import TorchMemorySaverAdapter
from sglang.srt.utils import (enable_show_time_cost,get_available_gpu_memory,init_custom_process_group,is_cuda,is_hip,monkey_patch_p2p_access_check,monkey_patch_vllm_gguf_config,set_cpu_offload_max_bytes,set_cuda_arch,
)logger = logging.getLogger(__name__)class ModelRunner:"""ModelRunner运行模型的前向传播过程，负责推理的核心工作。"""def __init__(self,model_config: ModelConfig,mem_fraction_static: float,gpu_id: int,tp_rank: int,tp_size: int,nccl_port: int,server_args: ServerArgs,is_draft_worker: bool = False,):# 解析传入的配置参数self.model_config = model_configself.mem_fraction_static = mem_fraction_staticself.device = server_args.deviceself.gpu_id = gpu_idself.tp_rank = tp_rankself.tp_size = tp_sizeself.dist_port = nccl_portself.server_args = server_argsself.is_draft_worker = is_draft_workerself.is_generation = model_config.is_generation  # 是否为生成模型self.is_multimodal = model_config.is_multimodal  # 是否为多模态模型self.should_log = tp_rank == 0  # 只有rank为0的节点进行日志记录self.spec_algorithm = SpeculativeAlgorithm.from_string(server_args.speculative_algorithm)# 针对特定模型的调整if (self.model_config.attention_arch == AttentionArch.MLAand not self.server_args.disable_mla):# MLA优化，仅适用于非CPU设备if self.server_args.device != "cpu":if server_args.enable_flashinfer_mla:logger.info("MLA优化已开启，使用flashinfer mla后端。")self.server_args.attention_backend = "flashinfer_mla"else:logger.info("MLA优化已开启，使用triton后端。")self.server_args.attention_backend = "triton"if self.server_args.enable_double_sparsity:# 启用双重稀疏优化，使用triton后端logger.info("双重稀疏优化已开启，使用triton后端，不启用CUDA图。")self.server_args.attention_backend = "triton"self.server_args.disable_cuda_graph = Trueif self.server_args.ds_heavy_channel_type is None:raise ValueError("请指定用于双重稀疏优化的重通道类型。")self.init_double_sparsity_channel_config(self.server_args.ds_heavy_channel_type)if self.is_multimodal:self.mem_fraction_static *= 0.95logger.info(f"由于这是一个多模态模型，自动将--mem-fraction-static减少为 {self.mem_fraction_static:.3f} ")# 特定架构的调整if self.model_config.hf_config.architectures == ["MllamaForConditionalGeneration"]:logger.info("为mllama模型自动关闭--chunked-prefill-size。")server_args.chunked_prefill_size = -1if self.model_config.hf_config.architectures == ["Qwen2VLForConditionalGeneration"]:# TODO: qwen2-vl暂时不支持基数缓存，自动禁用radix缓存logger.info("为qwen2-vl自动关闭--chunked-prefill-size并禁用基数缓存。")server_args.chunked_prefill_size = -1server_args.disable_radix_cache = True# 全局变量更新if server_args.show_time_cost:enable_show_time_cost()if server_args.disable_outlines_disk_cache:from outlines.caching import disable_cachedisable_cache()global_server_args_dict.update({"attention_backend": server_args.attention_backend,"sampling_backend": server_args.sampling_backend,"triton_attention_reduce_in_fp32": server_args.triton_attention_reduce_in_fp32,"disable_mla": server_args.disable_mla,"torchao_config": server_args.torchao_config,"enable_nan_detection": server_args.enable_nan_detection,"enable_dp_attention": server_args.enable_dp_attention,"enable_ep_moe": server_args.enable_ep_moe,"device": server_args.device,"enable_flashinfer_mla": server_args.enable_flashinfer_mla,"disable_radix_cache": server_args.disable_radix_cache,})set_cpu_offload_max_bytes(int(server_args.cpu_offload_gb * 1024**3))# 获取加载前的内存min_per_gpu_memory = self.init_torch_distributed()self.memory_saver_adapter = TorchMemorySaverAdapter.create(enable=self.server_args.enable_memory_saver)# 加载模型self.sampler = Sampler()self.load_model()# 应用torchao量化torchao_applied = getattr(self.model, "torchao_applied", False)# 在分层加载中，torchao可能已经被应用if not torchao_applied:apply_torchao_config_to_model(self.model, global_server_args_dict["torchao_config"])# 应用torch TP（如果模型支持）supports_torch_tp = getattr(self.model, "supports_torch_tp", False)if self.tp_size > 1 and supports_torch_tp:self.apply_torch_tp()self.torch_tp_applied = Trueelse:self.torch_tp_applied = False# 初始化内存池和注意力后端if server_args.lora_paths is not None:self.init_lora_manager()self.init_memory_pool(min_per_gpu_memory,server_args.max_running_requests,server_args.max_total_tokens,)if self.device == "cuda":self.init_cublas()self.init_attention_backend()self.init_cuda_graphs()else:self.cuda_graph_runner = Noneself.init_attention_backend()def init_torch_distributed(self):"""初始化torch分布式环境，包括通信后端和分布式训练设置。"""logger.info("初始化torch分布式环境开始。")torch.get_device_module(self.device).set_device(self.gpu_id)if self.device == "cuda":backend = "nccl"elif self.device == "xpu":backend = "gloo"elif self.device == "hpu":backend = "hccl"elif self.device == "cpu":backend = "gloo"if not self.server_args.enable_p2p_check:monkey_patch_p2p_access_check()# 初始化分布式环境if self.server_args.dist_init_addr:dist_init_method = f"tcp://{self.server_args.dist_init_addr}"else:dist_init_method = f"tcp://127.0.0.1:{self.dist_port}"set_custom_all_reduce(not self.server_args.disable_custom_all_reduce)if not self.is_draft_worker:init_distributed_environment(backend=backend,world_size=self.tp_size,rank=self.tp_rank,local_rank=self.gpu_id,distributed_init_method=dist_init_method,timeout=self.server_args.dist_timeout,)initialize_model_parallel(tensor_model_parallel_size=self.tp_size)initialize_dp_attention(enable_dp_attention=self.server_args.enable_dp_attention,tp_rank=self.tp_rank,tp_size=self.tp_size,dp_size=self.server_args.dp_size,)min_per_gpu_memory = get_available_gpu_memory(self.device, self.gpu_id, distributed=self.tp_size > 1)self.tp_group = get_tp_group()self.attention_tp_group = get_attention_tp_group()# 检查内存是否均衡if self.tp_size > 1:local_gpu_memory = get_available_gpu_memory(self.device, self.gpu_id)if min_per_gpu_memory < local_gpu_memory * 0.9:raise ValueError("内存容量不均衡，部分GPU可能被其他进程占用。")return min_per_gpu_memorydef load_model(self):"""加载模型并设置相关配置。"""logger.info(f"开始加载权重，当前内存={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB")# 设置线程数以减少线程冲突，提升加载速度if self.device != "cpu":torch.set_num_threads(1)if self.device == "cuda":if torch.cuda.get_device_capability()[0] < 8:logger.info("计算能力低于sm80，使用float16代替bfloat16。")self.server_args.dtype = "float16"self.model_config.dtype = torch.float16if torch.cuda.get_device_capability()[1] < 5:raise RuntimeError("SGLang仅支持sm75及以上。")set_cuda_arch()# 准备加载配置self.load_config = LoadConfig(load_format=self.server_args.load_format,download_dir=self.server_args.download_dir,)if self.server_args.load_format == "gguf":monkey_patch_vllm_gguf_config()# 加载模型monkey_patch_vllm_parallel_state()with self.memory_saver_adapter.region():self.model = get_model(model_config=self.model_config,load_config=self.load_config,device_config=DeviceConfig(self.device),)monkey_patch_vllm_parallel_state(reverse=True)# 加载KV缓存的缩放因子if self.server_args.kv_cache_dtype == "fp8_e4m3":if self.server_args.quantization_param_path is not None:if callable(getattr(self.model, "load_kv_cache_scales", None)):self.model.load_kv_cache_scales(self.server_args.quantization_param_path)logger.info("从%s加载KV缓存缩放因子",self.server_args.quantization_param_path,)else:raise RuntimeError("使用FP8 KV缓存并提供缩放因子，但模型不支持加载这些因子。")else:logger.warning("使用FP8 KV缓存，但没有提供缩放因子，默认使用1.0的因子。")# 设置其他参数self.sliding_window_size = (self.model.get_attention_sliding_window_size()if hasattr(self.model, "get_attention_sliding_window_size")else None)self.dtype = self.model_config.dtypelogger.info(f"权重加载结束，模型类型={type(self.model).__name__}, dtype={self.dtype}, "f"当前内存={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB")def update_weights_from_disk(self, model_path: str, load_format: str) -> tuple[bool, str]:"""从磁盘更新引擎权重。"""from sglang.srt.model_loader.loader import (DefaultModelLoader,device_loading_context,get_model_loader,)from sglang.srt.model_loader.utils import set_default_torch_dtypelogger.info(f"开始在线从磁盘更新引擎权重，当前内存={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB")target_device = torch.device(self.device)self.model_config.model_path = model_pathload_config = LoadConfig(load_format=load_format)# 当前仅支持vllm的DefaultModelLoaderloader = get_model_loader(load_config)if not isinstance(loader, DefaultModelLoader):message = f"获取模型加载器失败: {loader}."return False, messagedef get_weight_iter(config):iter = loader._get_weights_iterator(DefaultModelLoader.Source(config.model_path,revision=config.revision,fall_back_to_pt=getattr(self.model, "fall_back_to_pt_during_load", True),))return iterdef model_load_weights(model, iter):model.load_weights(iter)for _, module in self.model.named_modules():quant_method = getattr(module, "quant_method", None)if quant_method is not None:with device_loading_context(module, target_device):quant_method.process_weights_after_loading(module)return modelwith set_default_torch_dtype(self.model_config.dtype):try:iter = get_weight_iter(self.model_config)except Exception as e:message = f"获取权重迭代器失败: {e}."return False, messagetry:model = model_load_weights(self.model, iter)except Exception as e:message = (f"更新权重失败: {e}.\n将回滚到原始权重。")del itergc.collect()iter = get_weight_iter(self.model_config)self.model = model_load_weights(self.model, iter)return False, messageself.model = modelself.server_args.model_path = model_pathself.server_args.load_format = load_formatself.load_config = load_configlogger.info("更新权重结束。")return True, "成功从磁盘更新模型权重。"

TpModelWorkerClient类

"""A tensor parallel worker."""import dataclasses
import logging
import signal
import threading
from queue import Queue
from typing import Optionalimport psutil
import torchfrom sglang.srt.managers.io_struct import (GetWeightsByNameReqInput,InitWeightsUpdateGroupReqInput,UpdateWeightFromDiskReqInput,UpdateWeightsFromDistributedReqInput,UpdateWeightsFromTensorReqInput,
)
from sglang.srt.managers.schedule_batch import ModelWorkerBatch
from sglang.srt.managers.tp_worker import TpModelWorker
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import get_compiler_backend
from sglang.utils import get_exception_tracebacklogger = logging.getLogger(__name__)@torch.compile(dynamic=True, backend=get_compiler_backend())
def resolve_future_token_ids(input_ids, future_token_ids_map):"""解决未来的token id，通过查找映射更新输入的input_ids"""input_ids[:] = torch.where(input_ids < 0,future_token_ids_map[torch.clamp(-input_ids, min=0)],input_ids,)class TpModelWorkerClient:"""一个张量并行模型工作者客户端，负责与模型进行交互"""def __init__(self,server_args: ServerArgs,gpu_id: int,tp_rank: int,dp_rank: Optional[int],nccl_port: int,):# 初始化模型工作者self.worker = TpModelWorker(server_args, gpu_id, tp_rank, dp_rank, nccl_port)self.max_running_requests = self.worker.max_running_requestsself.device = self.worker.deviceself.gpu_id = gpu_id# 初始化未来token映射self.future_token_ids_ct = 0self.future_token_ids_limit = self.max_running_requests * 3self.future_token_ids_map = torch.empty((self.max_running_requests * 5,), dtype=torch.int32, device=self.device)# 启动线程self.input_queue = Queue()  # 输入队列，存储待处理的batchself.output_queue = Queue()  # 输出队列，存储处理结果self.forward_stream = torch.get_device_module(self.device).Stream()  # 为前向计算创建流self.forward_thread = threading.Thread(target=self.forward_thread_func,  # 启动前向线程)self.forward_thread.start()self.parent_process = psutil.Process().parent()  # 获取父进程self.scheduler_stream = torch.get_device_module(self.device).current_stream()if self.device == "cpu":self.scheduler_stream.synchronize = lambda: None  # 如果是CPU设备，则同步函数为空操作def get_worker_info(self):"""获取工作者信息，包括最大token数、内存池大小等"""return self.worker.get_worker_info()def get_pad_input_ids_func(self):"""获取填充input_ids的函数"""return self.worker.get_pad_input_ids_func()def get_tp_cpu_group(self):"""获取张量并行的CPU组"""return self.worker.get_tp_cpu_group()def get_attention_tp_cpu_group(self):"""获取注意力张量并行的CPU组"""return self.worker.get_attention_tp_cpu_group()def get_memory_pool(self):"""获取内存池，包含tokens和key-value对池"""return (self.worker.model_runner.req_to_token_pool,self.worker.model_runner.token_to_kv_pool,)def forward_thread_func(self):"""前向计算线程的函数，处理模型推理"""try:with torch.get_device_module(self.device).stream(self.forward_stream):self.forward_thread_func_()except Exception:traceback = get_exception_traceback()  # 获取异常堆栈logger.error(f"TpModelWorkerClient遇到异常: {traceback}")self.parent_process.send_signal(signal.SIGQUIT)  # 向父进程发送退出信号@torch.no_grad()def forward_thread_func_(self):"""前向计算的实际函数，处理模型的推理和结果收集"""batch_pt = 0batch_lists = [None] * 2while True:# 从输入队列获取batchmodel_worker_batch, future_token_ids_ct = self.input_queue.get()if not model_worker_batch:break  # 如果batch为空，退出循环# 保持对model_worker_batch的引用，否则它的tensor成员会被释放，导致CUDA内存访问错误batch_lists[batch_pt % 2] = model_worker_batchbatch_pt += 1# 创建事件self.launch_done = threading.Event()copy_done = torch.get_device_module(self.device).Event()# 解决输入中的未来token idsinput_ids = model_worker_batch.input_idsresolve_future_token_ids(input_ids, self.future_token_ids_map)# 进行前向推理logits_output, next_token_ids = self.worker.forward_batch_generation(model_worker_batch, self.launch_done)# 更新未来token id映射bs = len(model_worker_batch.seq_lens)self.future_token_ids_map[future_token_ids_ct + 1 : future_token_ids_ct + bs + 1] = next_token_ids# 将结果复制到CPUif model_worker_batch.return_logprob:logits_output.next_token_logprobs = (logits_output.next_token_logprobs.to("cpu", non_blocking=True))if logits_output.input_token_logprobs is not None:logits_output.input_token_logprobs = (logits_output.input_token_logprobs.to("cpu", non_blocking=True))if logits_output.hidden_states is not None:logits_output.hidden_states = logits_output.hidden_states.to("cpu", non_blocking=True)next_token_ids = next_token_ids.to("cpu", non_blocking=True)copy_done.record()# 将处理结果放入输出队列self.output_queue.put((copy_done, logits_output, next_token_ids))def resolve_batch_result(self, bid: int):"""从输出队列获取并处理一个batch的结果"""copy_done, logits_output, next_token_ids = self.output_queue.get()copy_done.synchronize()  # 等待复制完成self.launch_done.wait()  # 等待前向推理完成# 将logits结果转换为list格式if logits_output.next_token_logprobs is not None:logits_output.next_token_logprobs = (logits_output.next_token_logprobs.tolist())if logits_output.input_token_logprobs is not None:logits_output.input_token_logprobs = (logits_output.input_token_logprobs.tolist())next_token_ids = next_token_ids.tolist()  # 将token id转换为listreturn logits_output, next_token_idsdef forward_batch_generation(self, model_worker_batch: ModelWorkerBatch):"""生成新的batch并将其推送到队列进行处理"""# 创建sampling_info的副本，因为调度器会就地更新它sampling_info = model_worker_batch.sampling_infosampling_info.update_penalties()  # 更新惩罚model_worker_batch.sampling_info = self.cur_sampling_info = dataclasses.replace(sampling_info,sampling_info_done=threading.Event(),scaling_penalties=sampling_info.scaling_penalties,linear_penalties=sampling_info.linear_penalties,)# 同步CUDA流，避免非法内存访问错误self.scheduler_stream.synchronize()# 将新batch推送到输入队列self.input_queue.put((model_worker_batch, self.future_token_ids_ct))# 分配输出的future对象bs = len(model_worker_batch.seq_lens)future_next_token_ids = torch.arange(-(self.future_token_ids_ct + 1),-(self.future_token_ids_ct + 1 + bs),-1,dtype=torch.int32,device=self.device,)self.future_token_ids_ct = (self.future_token_ids_ct + bs) % self.future_token_ids_limitreturn None, future_next_token_idsdef update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput):"""从磁盘更新权重"""success, message = self.worker.update_weights_from_disk(recv_req)return success, messagedef init_weights_update_group(self, recv_req: InitWeightsUpdateGroupReqInput):"""初始化权重更新组，用于分布式训练"""success, message = self.worker.init_weights_update_group(recv_req)return success, messagedef update_weights_from_distributed(self, recv_req: UpdateWeightsFromDistributedReqInput):"""从分布式系统更新权重"""success, message = self.worker.update_weights_from_distributed(recv_req)return success, messagedef update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):"""从tensor更新权重"""success, message = self.worker.update_weights_from_tensor(recv_req)return success, messagedef get_weights_by_name(self, recv_req: GetWeightsByNameReqInput):"""根据权重名称获取模型权重"""return self.worker.get_weights_by_name(recv_req)def __delete__(self):"""析构函数，清理资源"""self.input_queue.put((None, None))  # 停止输入队列的处理self.copy_queue.put((None, None, None))  # 停止复制队列的处理

TpModelWorkerClient 类是一个负责管理张量并行模型工作者的客户端。它与 TpModelWorker 进行交互，管理输入输出队列并处理模型推理任务。
该类通过多线程实现前向推理，并在计算完成后将结果传回主线程。
各个方法实现了模型加载、前向推理、权重更新等功能。

Scheduler类

Scheduler 类及相关函数概述

1. Scheduler 类简介

Scheduler 类是一个调度器，负责管理和调度一个张量并行（TP）GPU工作者，它协调了模型的生成（generation）或嵌入（embedding）任务的执行。调度器通过接收请求、运行批次、并生成结果来处理请求，并管理与硬件（如GPU、内存池）相关的操作。它还涉及到缓存管理、内存池释放、权重更新等操作。

2. 主要职责与功能

• 调度批次：根据请求的优先级和资源的可用性，决定下一个需要执行的批次。
• 处理生成与嵌入请求：根据请求类型（生成或嵌入）调度对应的任务，处理模型生成的输出（如logits、token、embedding等）。
• 内存与缓存管理：通过内存池和缓存的管理，优化GPU资源的利用，避免内存泄漏。
• 异步处理：调度器支持异步任务处理，能够并行进行数据预填充与生成解码工作，提升吞吐量。
• Speculative Decoding（投机解码）：如果启用了投机解码算法，调度器会启动额外的草案工作者来提前生成部分预测，帮助加速推理过程。

3. 构造函数 (__init__)

• 初始化调度器，设置多种参数，包括模型配置、硬件信息（如GPU ID、张量并行大小等），以及缓存、内存池等。
• 根据配置，选择是否启用重叠调度（enable_overlap）。
• 初始化与其他模块的通信（如与tokenizer的IPC通信）。
• 启动与GPU工作者相关的实例（如 TpWorkerClass），以及用于投机解码的工作者（如 EAGLEWorker）。

4. 调度循环 (event_loop_normal, event_loop_overlap)

• event_loop_normal：正常的调度循环，依次接收请求、处理请求、执行批次并生成结果。
• event_loop_overlap：异步调度循环，能够重叠执行CPU处理与GPU计算，提升效率。

5. 批次相关函数

• get_next_batch_to_run：获取下一个要运行的批次，决定是否执行预填充（prefill）或解码（decode），并处理DP注意力机制。
• get_new_batch_prefill：检查是否有请求可以进行预填充，并根据策略生成新的批次。
• update_running_batch：更新当前正在运行的解码批次，检查是否内存溢出，必要时进行回退。
• run_batch：运行批次，调用相应的生成或嵌入函数，并返回结果。

6. 缓存和内存管理

• flush_cache：清空缓存和内存池，用于释放资源。
• check_memory：检查内存使用情况，发现内存泄漏时触发警告。

7. 请求处理

• recv_requests：接收来自不同源的请求，并将请求广播到其他张量并行工作者。
• process_input_requests：处理接收到的请求，分发到合适的处理函数（如生成请求、嵌入请求等）。
• handle_generate_request：处理生成请求，创建新的生成任务。
• handle_embedding_request：处理嵌入请求，生成相应的嵌入结果。

8. 日志与统计

• log_prefill_stats：记录预填充批次的统计数据，如缓存命中率、使用的token数等。
• log_decode_stats：记录解码批次的统计数据，如生成的token数、生成吞吐量等。

9. 权重更新与会话管理

• update_weights_from_disk、update_weights_from_distributed、update_weights_from_tensor：更新模型的权重，包括从磁盘、分布式或张量输入更新。
• open_session、close_session：用于管理会话，允许为每个会话创建、关闭任务。

10. 投机解码和混合式预填充

• 投机解码（Speculative Decoding）：如果启用，调度器会启动草案工作者（draft_worker），提前生成一些token，优化推理速度。
• 混合式预填充（Mixed Chunked Prefill）：支持混合式预填充，使得可以更高效地进行批次的生成。

总结

Scheduler 类是一个高度集成的调度器，负责协调多个工作者、管理请求、调度批次并优化GPU和内存资源的利用。它能够灵活处理生成与嵌入任务，支持异步操作、投机解码、混合预填充等高级功能，并能够高效管理缓存、内存以及模型权重的更新。

DataParallelController类

DataParallelController 类及相关函数概述

1. DataParallelController 类简介

DataParallelController 类是一个管理数据并行工作者的控制器。它负责根据不同的负载均衡方法（如轮询调度和最短队列调度）将请求分发到多个数据并行（DP）工作者。控制器还负责初始化、管理与数据并行相关的调度程序进程，并进行模型加载、请求调度等任务。

2. 主要职责与功能

• 负载均衡调度：根据负载均衡策略（ROUND_ROBIN 或 SHORTEST_QUEUE）选择合适的工作者，并将请求分发给工作者。
• 初始化与启动调度器：初始化与启动数据并行调度器进程，为每个数据并行工作者创建独立的调度程序。
• 调度请求：根据当前负载情况（如调度方法）将请求分发给不同的工作者。
• 跨进程通信：通过 ZeroMQ 实现跨进程通信，接收请求并将请求转发给合适的工作者。
• 模型信息管理：管理模型的信息，例如最大token数、请求输入长度等。

3. 构造函数 (__init__)

• 初始化调度器，设置服务器和端口参数。
• 选择负载均衡调度方法（轮询调度或最短队列调度）。
• 初始化 ZeroMQ 上下文和 IPC（进程间通信）套接字。
• 启动数据并行工作者，分别为每个 dp_rank 启动调度器进程。
• 配置调度器与工作者的通信通道。

4. 工作者启动函数

• launch_dp_schedulers：为每个数据并行工作者启动调度程序进程，并初始化相关的端口。
• launch_dp_attention_schedulers：为数据并行注意力机制的工作者启动调度程序进程，并计算注意力并行信息。
• launch_tensor_parallel_group：为每个Tensor并行组（TP group）启动进程，并设置工作者和调度器的通信端口。

5. 负载均衡调度方法

• round_robin_scheduler：按照轮询的方式将请求依次发送给每个工作者，循环利用所有工作者。
• shortest_queue_scheduler：通过最短队列的方式，将请求分发给当前队列最短的工作者（此方法尚未实现）。

6. 事件循环与请求调度

• event_loop：控制器的事件循环，持续接收请求，并根据负载均衡方法将请求分发给合适的工作者。如果是控制请求，则将请求转发给Tensor并行组的第一个工作者。

7. 辅助函数

• run_data_parallel_controller_process：启动 DataParallelController 进程，并处理日志记录和异常处理。
• setproctitle.setproctitle：设置进程标题，便于进程管理。
• configure_logger：配置日志记录。
• get_zmq_socket：创建 ZeroMQ 套接字，用于进程间通信。

8. 异常处理

• 在调度过程中，如果发生异常，会捕获并记录异常信息，并向父进程发送信号以终止程序。

总结

DataParallelController 类是一个用于管理数据并行工作者的控制器。它负责根据不同的负载均衡策略（如轮询或最短队列）调度请求到各个工作者，并进行进程间通信。它还负责启动和管理调度程序进程、模型信息的管理、内存池的释放等。该类使用 ZeroMQ 进行进程间通信，并通过多个子进程来实现数据并行的工作调度。

TokenizerManager 类

TokenizerManager 类是一个处理文本标记化的核心组件，负责管理与模型相关的各种功能，包括文本的标记化、批处理请求、会话管理、模型权重更新、性能监控等。以下是该类及其相关函数的总结：

主要功能组件：

1. 初始化 (__init__):

   • 初始化过程中，接受 server_args 和 port_args 参数，配置 tokenizer 管理器。
   • 通过 ZeroMQ（zmq）套接字实现进程间通信，分别用于从 detokenizer 接收数据和向调度器发送请求。
   • 加载模型配置并根据是否多模态（multimodal）加载相应的处理器和标记器。

2. 标记化处理:

   • _tokenize_one_request：处理单个请求的标记化，支持通过文本或嵌入向量（embeds）作为输入，并根据模型的上下文限制对输入进行检查。
   • generate_request：处理生成请求，标记化输入并将其发送给调度器进行处理。支持单个请求或批量请求的处理。

3. 会话管理:

   • open_session & close_session：负责创建和关闭会话，生成会话 ID 或关闭现有会话。
   • _handle_open_session_req_output：处理会话打开请求的响应。

4. 模型权重更新:

   • update_weights_from_disk, update_weights_from_distributed, update_weights_from_tensor：支持从磁盘、分布式系统或直接通过张量更新模型权重。
   • _wait_for_model_update_from_disk：等待并处理来自磁盘的模型更新请求，支持单节点或分布式环境中的更新。

5. 批处理请求:

   • _handle_batch_request：处理批量请求的标记化和处理，支持并发处理多个请求，并生成多个响应。它也处理并行采样和批量策略。
   • _handle_batch_output：处理批量输出，支持处理不同类型的结果，如字符串、token ID 和嵌入（embedding）等。

6. 请求管理:

   • _send_one_request：发送标记化的请求给调度器进行处理。
   • _wait_one_response：等待并处理单个请求的响应，支持超时或连接断开的处理。
   • create_abort_task：创建后台任务，用于在客户端断开连接时终止请求。

7. 内存和性能监控:

   • flush_cache, abort_request, start_profile, stop_profile：这些方法用于管理缓存刷新、请求中止以及性能分析控制。
   • release_memory_occupation, resume_memory_occupation：用于系统中的内存管理。

8. 日志和度量监控:

   • configure_logging：配置请求日志参数，如请求日志记录、请求转储文件夹等。
   • collect_metrics：收集性能度量指标，如第一次标记化时间、每个输出 token 的时间、端到端延迟等。
   • dump_requests：在达到设定的阈值时，将请求数据转储到文件中，以便离线分析。

9. 信号处理与优雅关闭:

   • auto_create_handle_loop：自动创建事件循环以处理信号和请求。
   • sigterm_watchdog：确保在关闭时优雅地处理未完成的请求，并最终结束进程。

辅助功能：

• detokenize_logprob_tokens：将 token 概率值和 token 索引反向标记化回文本或 token ID。
• detokenize_top_logprobs_tokens：与 detokenize_logprob_tokens 类似，但用于处理 top-k tokens 的标记化。
• convert_logprob_style：转换输出的 logprob 样式，包括返回 top-k logprobs 和其他相关的元数据。

作用：

TokenizerManager 类作为中间层，处理文本标记化和输入处理任务，管理请求的标记化、批处理请求、会话管理、模型更新、响应处理等。它与调度器、detokenizer 等其他系统组件进行交互，并提供了高效的请求处理和性能监控机制。

总结：

• 进程间通信通过 ZeroMQ 套接字进行。
• 模型标记化和生成标记化请求。
• 批量处理高效处理多个请求。
• 会话管理和请求状态追踪。
• 模型权重更新与同步。
• 日志记录和度量监控。
• 内存管理与性能分析。
• 信号处理与优雅关闭，确保进程在关闭时处理所有请求。

该类的核心目的是优化标记化任务、请求处理和模型管理，提供详细的性能监控和日志记录，以便进行高效的推理操作。

DetokenizerManager 类

DetokenizerManager 类及其相关函数的实现总结：

DetokenizerManager 是一个处理 解码 和 反tokenization（反向token化） 的类，它的主要功能是将模型的输出 token id 转换为可读文本，通常用于生成模型的输出。它作为一个独立进程运行，并通过与其他进程（如 tokenizer）进行消息交换来完成任务。它支持批量处理（即同时处理多个请求）。

1. __init__ 方法

• 初始化 DetokenizerManager，设置各种必要的参数和变量。
• 初始化 ZeroMQ 套接字，进行进程间通信：
  • recv_from_scheduler：从调度器接收任务（例如需要反token化的 token id）。
  • send_to_tokenizer：将解码结果发送给 tokenizer 进程。
• 根据传入的配置参数（如 skip_tokenizer_init）来决定是否初始化 tokenizer（即模型的反token化器）。
• 使用 TypeBasedDispatcher 来处理不同类型的请求：
  • BatchEmbeddingOut：无需解码，直接返回。
  • BatchTokenIDOut：需要解码。

2. trim_matched_stop 方法

• 用于处理解码过程中匹配到的停止符号（如 stop token 或文本），根据停止标志（finished_reason）来裁剪解码输出。
• 如果匹配到停止符号，裁剪掉停止符号后的部分。

3. event_loop 方法

• 这是主事件循环，负责不断接收调度器发送过来的任务，并将处理结果通过 ZeroMQ 套接字返回给 tokenizer 进程。
• 对接收到的任务（如 BatchTokenIDOut 类型的任务）进行处理，并调用相应的解码逻辑。

4. handle_batch_embedding_out 方法

• 处理 BatchEmbeddingOut 类型的输出。对于 embedding 类型的任务，不需要进行解码，直接返回原始的 embedding 输出。

5. handle_batch_token_id_out 方法

• 处理 BatchTokenIDOut 类型的输出。此方法的核心功能是 增量解码（incremental decoding）。
• 主要步骤：
  1. 遍历每个请求，检查 decode_status（解码状态），如果该请求的解码状态不存在或与当前解码版本不匹配，则初始化解码状态。
  2. 更新解码进度，进行 token ID 到文本的转换。
  3. 通过 tokenizer 反token化 token ids，将其转换为文本字符串。
  4. 对于流式解码的请求，根据是否结束来更新解码状态并生成部分输出文本。
  5. 合成最终的输出文本，并将结果封装为 BatchStrOut 类型发送回调度器。

6. LimitedCapacityDict 类

• 这是一个扩展自 OrderedDict 的自定义字典类，它用于存储解码状态（DecodeStatus）。
• 它有一个容量限制，当字典的大小超过设置的容量时，会移除最早加入的项，以确保字典不会占用过多内存。
• 通过这种方式，它能有效管理解码状态，避免内存溢出。

7. run_detokenizer_process 函数

• 该函数负责启动 DetokenizerManager 进程并处理异常。
• 它设置进程标题、日志配置，并初始化一个 DetokenizerManager 实例来运行解码任务。
• 如果遇到异常，会捕获并记录错误信息，并向父进程发送终止信号。

总结：

DetokenizerManager 是一个用于解码模型输出的管理器，它负责反token化 token id 并将其转换为文本。它与其他进程（如 tokenizer）进行通信，并能够处理大规模的批量解码任务。为了应对内存压力，它使用了带有限制容量的字典来管理解码状态，并通过增量解码的方式来逐步生成文本输出。此外，它还支持流式解码，允许在解码过程中逐步获取部分结果。