57: vLLM 核心模块逐文件：sampling.py

class SamplerBase:
    def __init__(self, vocab_size: int, device: torch.device):
        self.vocab_size = vocab_size
        self.device = device
    
    def sample(self, logits: torch.Tensor, **kwargs) -> torch.Tensor:
        """核心采样方法，将logits转换为token ids"""
        raise NotImplementedError("sample() must be implemented by subclasses")
    
    def preprocess_logits(self, logits: torch.Tensor, **kwargs) -> torch.Tensor:
        """logits预处理，如温度缩放、重复惩罚等"""
        return logits
    
    def postprocess_tokens(self, tokens: torch.Tensor, **kwargs) -> torch.Tensor:
        """token后处理，如约束检查、特殊token处理等"""
        return tokens

def create_sampling_pipeline(sampling_config: SamplingConfig):
    """创建采样流水线"""
    pipeline = []
    
    # 添加温度缩放组件
    if sampling_config.temperature > 0:
        pipeline.append(TemperatureScaler(sampling_config.temperature))
    
    # 添加Top-K过滤组件
    if sampling_config.top_k > 0:
        pipeline.append(TopKFilter(sampling_config.top_k))
    
    # 添加Top-P过滤组件
    if 0 < sampling_config.top_p < 1.0:
        pipeline.append(TopPFilter(sampling_config.top_p))
    
    # 添加重复惩罚组件
    if sampling_config.repetition_penalty != 1.0:
        pipeline.append(RepetitionPenalty(sampling_config.repetition_penalty))
    
    # 添加约束解码组件
    if sampling_config.constraints:
        pipeline.append(ConstraintDecoder(sampling_config.constraints))
    
    return pipeline

class TopKFilter:
    def __init__(self, top_k: int):
        self.top_k = top_k
    
    def __call__(self, logits: torch.Tensor) -> torch.Tensor:
        # 获取每个序列的top-k logits和对应的索引
        top_k_logits, top_k_indices = torch.topk(logits, self.top_k, dim=-1)
        
        # 创建一个掩码，将非top-k的logits设为负无穷
        mask = torch.full_like(logits, float('-inf'))
        mask.scatter_(-1, top_k_indices, top_k_logits)
        
        return mask

class TopPFilter:
    def __init__(self, top_p: float):
        self.top_p = top_p
    
    def __call__(self, logits: torch.Tensor) -> torch.Tensor:
        # 对logits进行softmax，得到概率分布
        probs = torch.softmax(logits, dim=-1)
        
        # 对概率进行排序（降序）
        sorted_probs, sorted_indices = torch.sort(probs, dim=-1, descending=True)
        
        # 计算累积概率
        cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
        
        # 创建掩码，保留累积概率<=top_p的Token
        # 同时确保至少保留一个Token
        mask = cumulative_probs <= self.top_p
        mask[..., 0] = True
        
        # 将掩码映射回原始索引
        mask = mask.scatter(1, sorted_indices.argsort(1), mask)
        
        # 应用掩码
        logits = logits.masked_fill(~mask, float('-inf'))
        
        return logits

class RepetitionPenalty:
    def __init__(self, penalty: float):
        self.penalty = penalty
    
    def __call__(self, logits: torch.Tensor, sequences: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, vocab_size = logits.shape
        
        # 遍历每个序列
        for i in range(batch_size):
            # 获取已生成的Token
            generated_tokens = sequences[i][:seq_len]
            
            # 统计每个Token的出现次数
            token_counts = torch.bincount(generated_tokens, minlength=vocab_size)
            
            # 对已出现的Token应用惩罚
            logits[i, :, token_counts > 0] /= self.penalty
        
        return logits

def vectorized_sample(logits: torch.Tensor, pipeline: List[SamplingComponent]) -> torch.Tensor:
    """向量化采样实现"""
    batch_size, vocab_size = logits.shape
    
    # 依次通过采样流水线
    processed_logits = logits
    for component in pipeline:
        processed_logits = component(processed_logits)
    
    # 对处理后的logits进行softmax
    probs = torch.softmax(processed_logits, dim=-1)
    
    # 使用Gumbel-Softmax技巧进行采样
    # 避免直接采样的不可微问题，同时保持高效
    gumbel_noise = -torch.log(-torch.log(torch.rand_like(probs) + 1e-10) + 1e-10)
    sampled_indices = torch.argmax(processed_logits + gumbel_noise, dim=-1)
    
    return sampled_indices

class ConstraintDecoder:
    def __init__(self, constraints: List[Dict]):
        self.constraints = constraints
        self.dfa_cache = {}
    
    def __call__(self, logits: torch.Tensor, current_state: Dict = None) -> torch.Tensor:
        # 根据当前状态和约束构建DFA
        dfa = self._build_dfa(current_state)
        
        # 获取允许的Token集合
        allowed_tokens = dfa.get_allowed_tokens()
        
        # 构建掩码，只允许生成符合约束的Token
        mask = torch.full_like(logits, float('-inf'))
        mask[:, allowed_tokens] = 0.0
        
        return logits + mask
    
    def _build_dfa(self, current_state: Dict) -> DFA:
        # 构建或从缓存中获取DFA
        # 实现细节省略...
        pass

from vllm import SamplingConfig

# 创建采样配置
sampling_config = SamplingConfig(
    temperature=0.7,
    top_p=0.95,
    top_k=50,
    repetition_penalty=1.1,
    max_tokens=512
)

from vllm import SamplingConfig

# 创建JSON约束
sampling_config = SamplingConfig(
    max_tokens=100,
    constraints=[
        {
            "type": "json",
            "schema": {
                "type": "object",
                "properties": {
                    "name": {"type": "string"},
                    "age": {"type": "integer"},
                    "city": {"type": "string"}
                },
                "required": ["name", "age", "city"]
            }
        }
    ]
)

from vllm.sampling import SamplerBase

class CustomSampler(SamplerBase):
    def __init__(self, vocab_size: int, device: torch.device, alpha: float):
        super().__init__(vocab_size, device)
        self.alpha = alpha
    
    def sample(self, logits: torch.Tensor) -> torch.Tensor:
        # 自定义采样逻辑
        # 示例：结合Top-K和Top-P的混合采样
        top_k = 50
        top_p = 0.95
        
        # Top-K过滤
        top_k_logits, top_k_indices = torch.topk(logits, top_k, dim=-1)
        
        # Top-P过滤
        probs = torch.softmax(top_k_logits, dim=-1)
        sorted_probs, sorted_indices = torch.sort(probs, dim=-1, descending=True)
        cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
        mask = cumulative_probs <= top_p
        mask[..., 0] = True
        
        # 应用掩码
        filtered_probs = probs * mask.gather(1, sorted_indices.argsort(1))
        filtered_probs = filtered_probs / filtered_probs.sum(dim=-1, keepdim=True)
        
        # 采样
        sampled_idx = torch.multinomial(filtered_probs, 1)
        sampled_token = top_k_indices.gather(1, sampled_idx)
        
        return sampled_token