import torch

def traditional_topk_sampling(logits, k=50):
    # 排序操作：O(V log V)，其中V是词汇表大小
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    
    # 取前K个token的概率
    top_k_logits = sorted_logits[:k]
    top_k_indices = sorted_indices[:k]
    
    # 重新归一化概率分布
    probs = torch.softmax(top_k_logits, dim=-1)
    
    # 从Top-K中采样
    sampled_idx = torch.multinomial(probs, num_samples=1)
    return top_k_indices[sampled_idx]

def traditional_topp_sampling(logits, p=0.9):
    # 排序操作：O(V log V)
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
    
    # 找到累积概率超过p的最小索引
    top_p_mask = cumulative_probs <= p
    top_p_logits = sorted_logits[top_p_mask]
    top_p_indices = sorted_indices[top_p_mask]
    
    # 归一化并采样
    probs = torch.softmax(top_p_logits, dim=-1)
    sampled_idx = torch.multinomial(probs, num_samples=1)
    return top_p_indices[sampled_idx]

import torch
import random

def inverse_transform_sampling(probs):
    """
    逆变换采样：直接从累积分布函数采样
    时间复杂度：O(V)，无需排序
    """
    # 计算累积分布
    cumsum_probs = torch.cumsum(probs, dim=-1)
    
    # 生成[0,1)之间的随机数
    u = torch.rand(1, device=probs.device)
    
    # 找到第一个大于等于u的位置
    sampled_idx = torch.searchsorted(cumsum_probs, u, right=True)
    return sampled_idx

def rejection_sampling_step(logits, max_prob=None):
    """
    单步拒绝采样
    """
    if max_prob is None:
        max_prob = torch.max(torch.softmax(logits, dim=-1))
    
    vocab_size = logits.size(-1)
    
    while True:
        # 从词汇表中随机选择一个位置
        rand_idx = torch.randint(0, vocab_size, (1,), device=logits.device)
        
        # 获取该位置的概率
        prob_at_idx = torch.softmax(logits, dim=-1)[rand_idx]
        
        # 从[0, max_prob]中采样
        u = torch.rand(1, device=logits.device) * max_prob
        
        # 拒绝采样判断
        if u <= prob_at_idx:
            return rand_idx  # 接受该样本
        # 否则继续循环，拒绝该样本

import torch
import flashinfer

# 准备logits（batch_size, vocab_size）
batch_size, vocab_size = 4, 100000
logits = torch.randn(batch_size, vocab_size, dtype=torch.float32, device="cuda")

# 使用FlashInfer的高效采样函数
from flashinfer.sampling import top_k_top_p_sampling_from_logits

# 批量采样
next_tokens = top_k_top_p_sampling_from_logits(
    logits,
    top_k=50,      # Top-K参数
    top_p=0.9,     # Top-P参数  
    temperature=0.7,  # 温度参数
    indices=None   # 可选：指定采样的位置索引
)

print(f"采样结果: {next_tokens}")

import torch
import flashinfer
from flashinfer.sampling import (
    top_k_top_p_sampling_from_logits,
    min_p_sampling_from_logits,
    sample_with_top_k_and_top_p_and_min_p
)

# 复合采样策略：同时使用Top-K、Top-P、Min-P
def advanced_sampling(logits, top_k=50, top_p=0.9, min_p=0.05, temperature=0.7):
    """
    高级采样：结合多种过滤策略
    """
    # 应用温度缩放
    scaled_logits = logits / temperature
    
    # 执行复合采样
    sampled_tokens = sample_with_top_k_and_top_p_and_min_p(
        scaled_logits,
        top_k=top_k,
        top_p=top_p,
        min_p=min_p
    )
    
    return sampled_tokens

# 性能优化的批量采样
def batch_sampling_optimized(logits_batch, sampling_params):
    """
    批量采样优化版本
    """
    batch_size = logits_batch.size(0)
    
    # 预分配输出tensor
    output_tokens = torch.empty(batch_size, dtype=torch.int64, device=logits_batch.device)
    
    # 分批处理以优化内存使用
    chunk_size = 64  # 根据GPU内存调整
    for i in range(0, batch_size, chunk_size):
        end_idx = min(i + chunk_size, batch_size)
        chunk_logits = logits_batch[i:end_idx]
        
        chunk_tokens = top_k_top_p_sampling_from_logits(
            chunk_logits,
            top_k=sampling_params['top_k'],
            top_p=sampling_params['top_p'],
            temperature=sampling_params['temperature']
        )
        
        output_tokens[i:end_idx] = chunk_tokens
    
    return output_tokens

# 示例使用
sampling_config = {
    'top_k': 50,
    'top_p': 0.9,
    'temperature': 0.7
}

# 创建模拟的推理批次
large_batch_logits = torch.randn(128, 100000, dtype=torch.float32, device="cuda")
sampled_results = batch_sampling_optimized(large_batch_logits, sampling_config)

# vLLM采样层的FlashInfer优化版本
class FlashInferSampler:
    def __init__(self):
        self.use_flashinfer = True  # 启用FlashInfer优化
    
    def forward(self, logits, sampling_metadata):
        """
        使用FlashInfer进行高效采样
        """
        if self.use_flashinfer and logits.is_cuda:
            # 使用FlashInfer的批量采样接口
            return self._flashinfer_sampling(logits, sampling_metadata)
        else:
            # 回退到传统采样方法
            return self._traditional_sampling(logits, sampling_metadata)
    
    def _flashinfer_sampling(self, logits, sampling_metadata):
        """
        FlashInfer优化的采样实现
        """
        from flashinfer.sampling import top_k_top_p_sampling_from_logits
        
        # 提取采样参数
        top_k_vals = []
        top_p_vals = []
        temp_vals = []
        
        for seq_group in sampling_metadata.seq_groups:
            sampling_params = seq_group.sampling_params
            top_k_vals.append(sampling_params.top_k)
            top_p_vals.append(sampling_params.top_p)
            temp_vals.append(sampling_params.temperature)
        
        # 转换为tensor
        top_k_tensor = torch.tensor(top_k_vals, device=logits.device)
        top_p_tensor = torch.tensor(top_p_vals, device=logits.device)
        temp_tensor = torch.tensor(temp_vals, device=logits.device)
        
        # 执行批量采样
        sampled_tokens = []
        for i, seq_group in enumerate(sampling_metadata.seq_groups):
            group_logits = logits[i:i+seq_group.num_seqs]
            
            # 对每个序列应用不同的采样参数
            group_tokens = top_k_top_p_sampling_from_logits(
                group_logits,
                top_k=top_k_vals[i],
                top_p=top_p_vals[i],
                temperature=temp_vals[i]
            )
            sampled_tokens.extend(group_tokens.tolist())
        
        return torch.tensor(sampled_tokens, device=logits.device)
    
    def _traditional_sampling(self, logits, sampling_metadata):
        """
        传统采样实现（用于对比和回退）
        """
        # 实现传统的排序采样逻辑
        # ...
        pass

import time
import torch
import numpy as np

def benchmark_sampling_methods():
    """
    对比传统采样与FlashInfer采样的性能
    """
    vocab_sizes = [50000, 100000, 150000]
    batch_sizes = [1, 4, 8, 16, 32]
    
    results = {}
    
    for vocab_size in vocab_sizes:
        print(f"\n=== 词汇表大小: {vocab_size} ===")
        results[vocab_size] = {}
        
        for batch_size in batch_sizes:
            print(f"批量大小: {batch_size}")
            
            # 生成测试数据
            logits = torch.randn(batch_size, vocab_size, dtype=torch.float32, device="cuda")
            
            # 预热GPU
            for _ in range(5):
                _ = traditional_topk_sampling(logits, k=50)
                _ = top_k_top_p_sampling_from_logits(
                    logits, top_k=50, top_p=0.9, temperature=0.7
                )
            
            # 传统采样性能测试
            torch.cuda.synchronize()
            start_event = torch.cuda.Event(enable_timing=True)
            end_event = torch.cuda.Event(enable_timing=True)
            
            start_event.record()
            for _ in range(100):  # 增加测试次数以获得更准确的结果
                traditional_result = traditional_topk_sampling(logits, k=50)
            end_event.record()
            torch.cuda.synchronize()
            traditional_time = start_event.elapsed_time(end_event) / 100 / 1000  # 转换为秒
            
            # FlashInfer采样性能测试
            start_event.record()
            for _ in range(100):
                flashinfer_result = top_k_top_p_sampling_from_logits(
                    logits, top_k=50, top_p=0.9, temperature=0.7
                )
            end_event.record()
            torch.cuda.synchronize()
            flashinfer_time = start_event.elapsed_time(end_event) / 100 / 1000
            
            # 计算加速比
            speedup = traditional_time / flashinfer_time
            
            results[vocab_size][batch_size] = {
                'traditional_time': traditional_time,
                'flashinfer_time': flashinfer_time,
                'speedup': speedup
            }
            
            print(f"  传统方法: {traditional_time:.4f}s")
            print(f"  FlashInfer: {flashinfer_time:.4f}s")  
            print(f"  加速比: {speedup:.2f}x")
    
    return results

# 运行基准测试
test_results = benchmark_sampling_methods()

在实际的基准测试中，我们还需要考虑更多的变量因素，如不同的temperature值、top_p值、以及不同类型的logits分布（例如，使用真实模型生成的logits vs 随机生成的logits）。这些因素都会影响采样算法的实际性能表现。

此外，为了确保测试结果的可靠性，建议在多个不同的硬件平台上重复测试，并使用统计方法（如置信区间）来评估结果的显著性。这对于生产环境中的技术选型决策尤为重要。

import torch.cuda as cuda

def analyze_gpu_utilization():
    """
    分析GPU资源利用率
    """
    # 监控GPU内存使用
    initial_memory = cuda.memory_allocated()
    
    # 执行采样操作
    logits = torch.randn(16, 100000, dtype=torch.float32, device="cuda")
    
    # 传统方法
    start_event = cuda.Event(enable_timing=True)
    end_event = cuda.Event(enable_timing=True)
    
    start_event.record()
    traditional_result = traditional_topk_sampling(logits, k=50)
    end_event.record()
    cuda.synchronize()
    
    traditional_time = start_event.elapsed_time(end_event)
    traditional_memory = cuda.memory_allocated() - initial_memory
    
    # FlashInfer方法
    start_event.record()
    flashinfer_result = top_k_top_p_sampling_from_logits(
        logits, top_k=50, top_p=0.9, temperature=0.7
    )
    end_event.record()
    cuda.synchronize()
    
    flashinfer_time = start_event.elapsed_time(end_event)
    flashinfer_memory = cuda.memory_allocated() - initial_memory
    
    print(f"GPU内存使用对比:")
    print(f"  传统方法: {traditional_memory / 1024**2:.2f} MB")
    print(f"  FlashInfer: {flashinfer_memory / 1024**2:.2f} MB")
    print(f"  内存节省: {(traditional_memory - flashinfer_memory) / traditional_memory * 100:.1f}%")

# vLLM配置文件示例
"""
vllm_config.yaml
"""
model: "meta-llama/Llama-2-7b-chat-hf"
tensor_parallel_size: 1
dtype: "float16"
enable_prefix_caching: true
# 启用FlashInfer采样优化
use_flashinfer_sampler: true

# 采样参数配置
sampling_params:
  top_k: 50
  top_p: 0.9
  temperature: 0.7
  min_p: 0.05

# SGLang运行时配置
import sglang as sgl

@sgl.function
def llm_generation(s, prompt: str):
    # 使用FlashInfer优化的采样
    s += sgl.user(prompt)
    s += sgl.assistant(
        sgl.gen("response", 
                max_tokens=512,
                temperature=0.7,
                top_p=0.9,
                top_k=50)
    )

# 启用FlashInfer后端
runtime = sgl.Runtime(
    model_path="meta-llama/Llama-2-7b-chat-hf",
    backend="flashinfer"  # 指定使用FlashInfer后端
)

# 执行生成
result = llm_generation.run(prompt="Explain quantum computing")

class SamplingPerformanceMonitor:
    def __init__(self):
        self.metrics = {
            'sampling_times': [],
            'acceptance_rates': [],
            'rejection_counts': []
        }
    
    def monitor_sampling_operation(self, logits, sampling_params):
        """
        监控采样操作的性能指标
        """
        import time
        
        start_time = time.time()
        
        # 执行采样并记录统计信息
        result = top_k_top_p_sampling_from_logits(
            logits, 
            top_k=sampling_params.get('top_k', 50),
            top_p=sampling_params.get('top_p', 0.9),
            temperature=sampling_params.get('temperature', 1.0)
        )
        
        sampling_time = time.time() - start_time
        self.metrics['sampling_times'].append(sampling_time)
        
        # 计算有效采样率（基于拒绝采样的统计）
        if hasattr(result, 'acceptance_rate'):
            self.metrics['acceptance_rates'].append(result.acceptance_rate)
        
        return result
    
    def get_performance_summary(self):
        """
        获取性能摘要
        """
        if not self.metrics['sampling_times']:
            return "No sampling operations recorded"
        
        avg_time = np.mean(self.metrics['sampling_times'])
        std_time = np.std(self.metrics['sampling_times'])
        
        summary = {
            'avg_sampling_time_ms': avg_time * 1000,
            'std_sampling_time_ms': std_time * 1000,
            'total_operations': len(self.metrics['sampling_times']),
            'q95_time_ms': np.percentile(self.metrics['sampling_times'], 95) * 1000
        }
        
        return summary