# vllm/v1/spec_decode/eagle.py

class SpecDecodeBaseProposer:
    def __init__(
        self,
        vllm_config: VllmConfig,
        device: torch.device,
        pass_hidden_states_to_model: bool,  # 是否传递隐藏状态
        runner=None,
    ):
        self.vllm_config = vllm_config
        self.speculative_config = vllm_config.speculative_config
        self.draft_model_config = self.speculative_config.draft_model_config
        self.method = self.speculative_config.method

        # 配置参数
        self.num_speculative_tokens = self.speculative_config.num_speculative_tokens
        self.hidden_size = self.draft_model_config.get_hidden_size()

        # 持久化缓冲区（用于 CUDA Graph）
        self.input_ids = torch.zeros(
            self.max_num_tokens, dtype=torch.int32, device=device
        )
        self.positions = torch.zeros(
            self.max_num_tokens, dtype=torch.int64, device=device
        )
        self.hidden_states = torch.zeros(
            (self.max_num_tokens, self.hidden_size),
            dtype=self.dtype, device=device
        )

def propose(
    self,
    target_token_ids: torch.Tensor,      # [num_tokens]
    target_positions: torch.Tensor,       # [num_tokens]
    target_hidden_states: torch.Tensor,   # [num_tokens, hidden_size]
    next_token_ids: torch.Tensor,         # [batch_size]
    last_token_indices: torch.Tensor | None,
    common_attn_metadata: CommonAttentionMetadata,
    sampling_metadata: SamplingMetadata,
    ...
) -> torch.Tensor:
    batch_size = common_attn_metadata.batch_size()

    # 1. 设置第一次前向传播的输入
    num_tokens, last_token_indices, common_attn_metadata = (
        self.set_inputs_first_pass(
            target_token_ids=target_token_ids,
            next_token_ids=next_token_ids,
            target_positions=target_positions,
            last_token_indices=last_token_indices,
            cad=common_attn_metadata,
            num_rejected_tokens_gpu=num_rejected_tokens_gpu,
        )
    )

    # 2. 构建 Attention Metadata
    attn_metadata = attn_metadata_builder.build_for_drafting(
        common_attn_metadata=common_attn_metadata, draft_index=0
    )

    # 3. 第一次前向传播
    with set_forward_context(per_layer_attn_metadata, ...):
        ret_hidden_states = self.model(
            input_ids=self.input_ids[:num_input_tokens],
            positions=self._get_positions(num_input_tokens),
            hidden_states=self.hidden_states[:num_input_tokens]
                          if self.pass_hidden_states_to_model else None,
        )

    # 4. 采样第一个 draft token
    logits = self.model.compute_logits(last_hidden_states[last_token_indices])
    draft_token_ids = logits.argmax(dim=-1)
    draft_token_ids_list = [draft_token_ids]

    # 5. 生成剩余的 draft tokens
    for token_index in range(self.num_speculative_tokens - 1):
        # 更新输入
        input_ids = draft_token_ids_list[-1].int()
        positions += 1

        # 处理超出最大长度的情况
        exceeds_max_model_len = positions >= self.max_model_len
        clamped_positions = torch.where(exceeds_max_model_len, 0, positions)

        # 更新序列长度
        common_attn_metadata.seq_lens += 1

        # 计算新的 slot mapping
        block_numbers = clamped_positions // block_size
        block_ids = common_attn_metadata.block_table_tensor.gather(
            dim=1, index=block_numbers.view(-1, 1)
        )
        common_attn_metadata.slot_mapping = (
            block_ids * block_size + clamped_positions % block_size
        )
        # 屏蔽超长位置
        common_attn_metadata.slot_mapping.masked_fill_(
            exceeds_max_model_len, PADDING_SLOT_ID
        )

        # 执行模型
        with set_forward_context(...):
            ret_hidden_states = self.model(
                input_ids=self.input_ids[:input_batch_size],
                positions=self._get_positions(input_batch_size),
                hidden_states=self.hidden_states[:input_batch_size],
            )

        # 采样下一个 token
        logits = self.model.compute_logits(last_hidden_states[:batch_size])
        draft_token_ids = logits.argmax(dim=-1)
        draft_token_ids_list.append(draft_token_ids)

    # 返回所有 draft tokens
    return torch.stack(draft_token_ids_list, dim=1)

# vllm/v1/spec_decode/draft_model.py

class DraftModelProposer(SpecDecodeBaseProposer):
    def __init__(
        self,
        vllm_config: VllmConfig,
        device: torch.device,
        runner=None,
    ):
        super().__init__(
            vllm_config=vllm_config,
            device=device,
            pass_hidden_states_to_model=False,  # 不需要 Target 的隐藏状态
            runner=runner,
        )
        # 验证约束
        self._raise_if_multimodal()     # 暂不支持多模态
        self._raise_if_mrope()          # 暂不支持 M-RoPE
        self._raise_if_vocab_size_mismatch()  # 词表大小必须一致
        self._raise_if_draft_tp_mismatch()    # TP 大小必须一致

def set_inputs_first_pass(
    self,
    target_token_ids: torch.Tensor,
    next_token_ids: torch.Tensor,
    target_positions: torch.Tensor,
    last_token_indices: torch.Tensor | None,
    cad: CommonAttentionMetadata,
    num_rejected_tokens_gpu: torch.Tensor | None,
) -> tuple[int, torch.Tensor, CommonAttentionMetadata]:
    batch_size = cad.batch_size()

    # 使用 Triton kernel 合并 tokens
    # target_toks: [a1, b1, b2, c1, c2, c3]
    # next_toks:   [a2, b3, c4]
    # 结果:        [a1, a2, b1, b2, b3, c1, c2, c3, c4]
    merge_toks_kernel[grid](
        target_toks_ptr=target_token_ids,
        next_toks_ptr=next_token_ids,
        query_start_locs_ptr=start_locs,
        query_end_locs_ptr=end_locs,
        out_ptr_merged_toks=self.input_ids,
        out_ptr_is_rejected_tok=is_rejected_tok,
        ...
    )

    # 重新计算 slot mapping
    new_slot_mapping = compute_new_slot_mapping(
        cad=cad,
        new_positions=self.positions[:num_tokens],
        is_rejected_token_mask=is_rejected_tok,
        block_size=self._block_size(),
        max_model_len=self.max_model_len,
    )

    return num_tokens, new_last_token_indices, new_cad

def load_model(self, target_model: Any) -> None:
    # 创建 Draft 模型专用的 VllmConfig
    draft_vllm_config = create_vllm_config_for_draft_model(
        target_model_vllm_config=self.vllm_config
    )

    logger.info(
        "Starting to load draft model %s. TP=%d, rank=%d",
        draft_vllm_config.model_config.model,
        draft_vllm_config.parallel_config.tensor_parallel_size,
        draft_vllm_config.parallel_config.rank,
    )

    # 加载模型并设置编译标签
    with set_model_tag("draft_model"):
        self.model = get_model(vllm_config=draft_vllm_config, prefix="draft_model")

# vllm/v1/spec_decode/eagle.py

class EagleProposer(SpecDecodeBaseProposer):
    def __init__(
        self,
        vllm_config: VllmConfig,
        device: torch.device,
        runner=None,
    ):
        super().__init__(
            vllm_config,
            device,
            pass_hidden_states_to_model=True,  # 需要传递隐藏状态
            runner=runner,
        )

def load_model(self, target_model: nn.Module) -> None:
    # 加载 EAGLE head
    with set_model_tag("eagle_head"):
        self.model = get_model(
            vllm_config=self.vllm_config,
            model_config=draft_model_config
        )

    # 检查是否需要共享 embedding
    if hasattr(self.model, "has_own_embed_tokens"):
        if not self.model.has_own_embed_tokens:
            share_embeddings = True
            logger.info("Sharing target model embedding weights with draft model")

    if share_embeddings:
        # 共享 embed_tokens
        del self.model.model.embed_tokens
        self.model.model.embed_tokens = target_embed_tokens

    # 共享 lm_head
    if share_lm_head:
        del self.model.lm_head
        self.model.lm_head = target_language_model.lm_head

# vllm/v1/worker/gpu/spec_decode/rejection_sample.py

@triton.jit
def _rejection_sample_kernel(
    sampled_ptr,           # [num_reqs, num_speculative_steps + 1]
    sampled_stride,
    num_sampled_ptr,       # [num_reqs]
    target_sampled_ptr,    # [num_draft_tokens + num_reqs]
    input_ids_ptr,         # [num_draft_tokens + num_reqs]（draft tokens）
    cu_num_logits_ptr,     # [num_reqs + 1]
):
    req_idx = tl.program_id(0)
    start_idx = tl.load(cu_num_logits_ptr + req_idx)
    end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
    num_tokens = end_idx - start_idx

    num_sampled = 0
    rejected = False

    # 逐个比较 draft token 和 target token
    for i in range(num_tokens - 1):
        if not rejected:
            target_sampled = tl.load(target_sampled_ptr + start_idx + i)
            draft_sampled = tl.load(input_ids_ptr + start_idx + i + 1)

            # 存储 target 的采样结果
            tl.store(sampled_ptr + req_idx * sampled_stride + i, target_sampled)
            num_sampled += 1

            # 检查是否匹配
            if target_sampled != draft_sampled:
                rejected = True  # 一旦不匹配，后续全部拒绝

    # 处理最后一个 token（bonus token）
    if not rejected:
        target_sampled = tl.load(target_sampled_ptr + start_idx + num_tokens - 1)
        tl.store(
            sampled_ptr + req_idx * sampled_stride + num_tokens - 1,
            target_sampled
        )
        num_sampled += 1

    tl.store(num_sampled_ptr + req_idx, num_sampled)

def propose_tree(
    self,
    batch_size: int,
    logits: torch.Tensor,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
    common_attn_metadata: CommonAttentionMetadata,
    ...
) -> list[torch.Tensor]:
    # 解析投机 token 树
    tree_depth = len(self.cu_drafts_per_level)

    # 第一层：从 root 采样多个候选
    num_children = self.child_drafts_per_level[0]
    if num_children == 1:
        draft_token_ids = logits.argmax(dim=-1).view(batch_size, -1)
    else:
        # Top-K 采样
        draft_token_ids = torch.topk(
            logits, num_children, dim=-1
        ).indices.view(batch_size, -1)

    draft_token_ids_list = [draft_token_ids]

    # 逐层生成
    for level in range(tree_depth - 1):
        # 构建树注意力元数据
        attn_metadata = tree_attn_metadata_builder.build_for_drafting(
            common_attn_metadata=common_attn_metadata,
            draft_index=level + 1
        )

        # 执行前向传播
        with set_forward_context(...):
            last_hidden_states, hidden_states = self.model(
                input_ids=self.input_ids[:num_input_tokens],
                positions=self.positions[:num_input_tokens],
                hidden_states=self.hidden_states[:num_input_tokens],
            )

        # 为下一层采样
        logits = self.model.compute_logits(draft_last_hidden_states)
        num_children = self.child_drafts_per_level[level + 1]
        if num_children == 1:
            draft_token_ids = logits.argmax(dim=-1).view(batch_size, -1)
        else:
            draft_token_ids = torch.topk(
                logits, num_children, dim=-1
            ).indices.view(batch_size, -1)

        draft_token_ids_list.append(draft_token_ids)

    return draft_token_ids_list

from vllm import LLM, SamplingParams

# 使用独立 Draft 模型
llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    speculative_model="meta-llama/Llama-3.2-1B-Instruct",
    num_speculative_tokens=5,  # 每次投机生成的 token 数
)

# 使用 EAGLE
llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    speculative_model="path/to/eagle-head",
    speculative_method="eagle",
    num_speculative_tokens=5,
)

# 使用 Draft 模型
vllm serve meta-llama/Llama-3.1-70B-Instruct \
    --speculative-model meta-llama/Llama-3.2-1B-Instruct \
    --num-speculative-tokens 5

# 使用 EAGLE
vllm serve meta-llama/Llama-3.1-70B-Instruct \
    --speculative-model path/to/eagle-head \
    --speculative-method eagle \
    --num-speculative-tokens 5

class SpecDecodeBaseProposer:
    def initialize_cudagraph_keys(self, cudagraph_mode: CUDAGraphMode) -> None:
        """初始化 CUDA Graph dispatcher"""
        if (
            not self.speculative_config.enforce_eager
            and cudagraph_mode.mixed_mode() in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]
        ):
            eagle_cudagraph_mode = CUDAGraphMode.PIECEWISE
        else:
            eagle_cudagraph_mode = CUDAGraphMode.NONE

        self.cudagraph_dispatcher.initialize_cudagraph_keys(eagle_cudagraph_mode)

def prepare_inputs_padded(
    self,
    common_attn_metadata: CommonAttentionMetadata,
    spec_decode_metadata: SpecDecodeMetadata,
    valid_sampled_tokens_count: torch.Tensor,
) -> tuple[CommonAttentionMetadata, torch.Tensor, torch.Tensor]:
    """
    准备投机解码的输入，使用 padding 保持固定形状。
    被拒绝的 token 作为 padding，后续通过 token_indices_to_sample 过滤。
    """
    num_reqs = common_attn_metadata.num_reqs
    device = valid_sampled_tokens_count.device

    token_indices_to_sample = torch.empty(
        (num_reqs,), dtype=torch.int32, device=device
    )
    num_rejected_tokens_gpu = torch.empty(
        (num_reqs,), dtype=torch.int32, device=device
    )

    # 使用 Triton kernel 计算
    eagle_prepare_inputs_padded_kernel[grid](
        spec_decode_metadata.cu_num_draft_tokens,
        valid_sampled_tokens_count,
        common_attn_metadata.query_start_loc,
        token_indices_to_sample,
        num_rejected_tokens_gpu,
        num_reqs,
    )

    return spec_common_attn_metadata, token_indices_to_sample, num_rejected_tokens_gpu

# vllm/model_executor/layers/quantization/__init__.py

QuantizationMethods = Literal[
    "awq",           # Activation-aware Weight Quantization
    "fp8",           # FP8 量化
    "gptq",          # Post-Training Quantization for GPT
    "gptq_marlin",   # GPTQ with Marlin kernel
    "awq_marlin",    # AWQ with Marlin kernel
    "bitsandbytes",  # BitsAndBytes 量化
    "gguf",          # GGUF 格式量化
    "compressed-tensors",  # 压缩张量
    "torchao",       # PyTorch AO 量化
    "modelopt",      # NVIDIA ModelOpt FP8
    "mxfp4",         # MXFP4 格式
    ...
]

# vllm/model_executor/layers/quantization/base_config.py

class QuantizationConfig(ABC):
    """量化配置基类"""

    @abstractmethod
    def get_name(self) -> QuantizationMethods:
        """量化方法名称"""
        raise NotImplementedError

    @abstractmethod
    def get_supported_act_dtypes(self) -> list[torch.dtype]:
        """支持的激活数据类型"""
        raise NotImplementedError

    @classmethod
    @abstractmethod
    def get_min_capability(cls) -> int:
        """最低 GPU 计算能力要求
        70 = Volta, 75 = Turing, 80 = Ampere
        """
        raise NotImplementedError

    @abstractmethod
    def get_quant_method(
        self, layer: torch.nn.Module, prefix: str
    ) -> QuantizeMethodBase | None:
        """获取量化方法实现"""
        raise NotImplementedError


class QuantizeMethodBase(ABC):
    """量化方法基类"""

    @abstractmethod
    def create_weights(
        self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs
    ):
        """创建量化权重"""
        raise NotImplementedError

    @abstractmethod
    def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor:
        """应用量化计算"""
        raise NotImplementedError

# vllm/model_executor/layers/quantization/fp8.py

class Fp8Config(QuantizationConfig):
    """FP8 量化配置"""

    def __init__(
        self,
        is_checkpoint_fp8_serialized: bool = False,  # 是否为 FP8 序列化的检查点
        activation_scheme: str = "dynamic",          # 激活量化方案
        ignored_layers: list[str] | None = None,     # 忽略的层
        weight_block_size: list[int] | None = None,  # 块大小
    ) -> None:
        super().__init__()
        self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized

        # 支持的激活方案: static, dynamic
        if activation_scheme not in ACTIVATION_SCHEMES:
            raise ValueError(f"Unsupported activation scheme {activation_scheme}")
        self.activation_scheme = activation_scheme
        self.ignored_layers = ignored_layers or []
        self.weight_block_size = weight_block_size

    @classmethod
    def get_min_capability(cls) -> int:
        return 75  # 最低支持 Turing 架构

    @classmethod
    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
        return [torch.bfloat16, torch.half]

# 加载 AWQ 量化模型
from vllm import LLM

llm = LLM(
    model="TheBloke/Llama-2-7B-AWQ",
    quantization="awq",
)

# 或者使用 Marlin 加速
llm = LLM(
    model="TheBloke/Llama-2-7B-AWQ",
    quantization="awq_marlin",  # 使用 Marlin 内核
)

# 加载 GPTQ 量化模型
llm = LLM(
    model="TheBloke/Llama-2-7B-GPTQ",
    quantization="gptq",
)

# 使用 Marlin 内核加速
llm = LLM(
    model="TheBloke/Llama-2-7B-GPTQ",
    quantization="gptq_marlin",
)

llm = LLM(
    model="meta-llama/Llama-2-7B",
    quantization="bitsandbytes",
    load_format="bitsandbytes",
)

llm = LLM(
    model="TheBloke/Llama-2-7B-GGUF",
    quantization="gguf",
)

# vllm/model_executor/layers/quantization/kv_cache.py

class BaseKVCacheMethod:
    """KV Cache 量化基类"""

    def quant_kv_tensor(
        self,
        key: torch.Tensor,
        value: torch.Tensor,
        scale: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """量化 KV 张量"""
        raise NotImplementedError

llm = LLM(
    model="meta-llama/Llama-2-7B",
    kv_cache_dtype="fp8",  # 或 "fp8_e4m3", "fp8_e5m2"
)

# vllm/model_executor/model_loader/online_quantization.py

# 动态 FP8 量化
llm = LLM(
    model="meta-llama/Llama-2-7B",
    quantization="fp8",
    # 将 FP16 权重动态转换为 FP8
)

# 高精度生产环境
llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    quantization="fp8",
    tensor_parallel_size=4,
)

# 资源受限环境
llm = LLM(
    model="TheBloke/Llama-2-70B-GPTQ",
    quantization="gptq_marlin",
    tensor_parallel_size=2,
)

# 单 GPU 部署
llm = LLM(
    model="TheBloke/Llama-2-7B-AWQ",
    quantization="awq_marlin",
    max_model_len=8192,
)

# vllm/distributed/parallel_state.py

class GroupCoordinator:
    """分布式组协调器"""

    def __init__(
        self,
        group_ranks: list[list[int]],
        local_rank: int,
        torch_distributed_backend: str,
        use_pynccl: bool,
        ...
    ):
        self.rank = torch.distributed.get_rank()
        self.ranks = group_ranks[local_rank]
        self.world_size = len(self.ranks)
        self.local_rank = local_rank

        # 创建通信组
        self.device_group = torch.distributed.new_group(
            self.ranks, backend=torch_distributed_backend
        )

    def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
        """AllReduce 操作"""
        if self.world_size == 1:
            return input_

        if self.use_custom_op_call:
            return torch.ops.vllm.all_reduce(input_, group_name=self.unique_name)
        else:
            return self._all_reduce_out_place(input_)

    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
        """AllGather 操作"""
        if self.world_size == 1:
            return input_

        if self.use_custom_op_call:
            return torch.ops.vllm.all_gather(
                input_, dim, self.world_size, group_name=self.unique_name
            )
        else:
            return self._all_gather_out_place(input_, dim)

    def reduce_scatter(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
        """ReduceScatter 操作"""
        if self.world_size == 1:
            return input_

        return self._reduce_scatter_out_place(input_, dim)

from vllm import LLM

# 配置流水线并行
llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    tensor_parallel_size=2,   # 每个流水线阶段 2 卡张量并行
    pipeline_parallel_size=2, # 2 个流水线阶段
    # 总共需要 2 × 2 = 4 张 GPU
)

# 启动多个 vLLM 实例
# 实例 1：使用 GPU 0-1
CUDA_VISIBLE_DEVICES=0,1 vllm serve model --tensor-parallel-size 2 --port 8000

# 实例 2：使用 GPU 2-3
CUDA_VISIBLE_DEVICES=2,3 vllm serve model --tensor-parallel-size 2 --port 8001

# vllm/distributed/device_communicators/pynccl.py

class PyNcclCommunicator:
    """PyNccl 通信器"""

    def __init__(self, group: ProcessGroup, device: torch.device):
        self.group = group
        self.device = device

        # 初始化 NCCL 通信
        self.nccl_comm = self._init_nccl()

    def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor:
        """使用 NCCL 执行 AllReduce"""
        # 调用 NCCL AllReduce
        return self._nccl_all_reduce(tensor)

# vllm/distributed/device_communicators/custom_all_reduce.py

class CustomAllReduce:
    """
    自定义 AllReduce，针对小张量优化。
    使用共享内存和 CUDA 内核实现低延迟通信。
    """

    def __init__(self, group: ProcessGroup):
        self.group = group
        # 分配共享内存用于通信
        self._init_shared_memory()

    def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor:
        # 对于小张量，使用自定义内核
        if tensor.numel() < self.threshold:
            return self._custom_all_reduce(tensor)
        # 对于大张量，使用 NCCL
        return self._nccl_all_reduce(tensor)

# 并行组划分示例
# 假设 4 GPU，TP=2，PP=2

# 张量并行组：
# [GPU 0, GPU 1]  # 第一阶段
# [GPU 2, GPU 3]  # 第二阶段

# 流水线并行组：
# [GPU 0, GPU 2]  # 第一个数据并行副本
# [GPU 1, GPU 3]  # 第二个数据并行副本

# vllm/v1/executor/multiproc_executor.py

class MultiprocExecutor:
    """多进程执行器，用于单节点多 GPU"""

    def __init__(self, vllm_config: VllmConfig):
        self.vllm_config = vllm_config
        parallel_config = vllm_config.parallel_config

        # 计算总 Worker 数
        self.world_size = (
            parallel_config.tensor_parallel_size *
            parallel_config.pipeline_parallel_size
        )

        # 启动 Worker 进程
        self.workers = self._start_workers()

    def _start_workers(self):
        """使用 multiprocessing 启动 Worker"""
        workers = []
        for rank in range(self.world_size):
            worker = multiprocessing.Process(
                target=self._worker_main,
                args=(rank,)
            )
            worker.start()
            workers.append(worker)
        return workers

# 使用 Ray 进行多节点分布式推理
from vllm import LLM

# Ray 会自动检测集群中的 GPU
llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    tensor_parallel_size=4,  # 使用 4 张 GPU
    distributed_executor_backend="ray",
)

# vllm/distributed/kv_transfer/kv_connector/v1/base.py

class KVConnectorBase:
    """KV Cache 传输连接器基类"""

    def send_kv_cache(
        self,
        request_id: str,
        kv_cache: torch.Tensor,
    ) -> None:
        """发送 KV Cache 到远程节点"""
        raise NotImplementedError

    def recv_kv_cache(
        self,
        request_id: str,
    ) -> torch.Tensor:
        """从远程节点接收 KV Cache"""
        raise NotImplementedError

# 单节点 4 GPU 张量并行
llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    tensor_parallel_size=4,
)

# 命令行方式
vllm serve meta-llama/Llama-3.1-70B-Instruct --tensor-parallel-size 4

# 4 路张量并行 + 2 路流水线并行
llm = LLM(
    model="meta-llama/Llama-3.1-405B-Instruct",
    tensor_parallel_size=4,
    pipeline_parallel_size=2,
)

# 节点 1（主节点）
vllm serve meta-llama/Llama-3.1-405B-Instruct \
    --tensor-parallel-size 8 \
    --pipeline-parallel-size 2 \
    --distributed-executor-backend ray

# 确保 Ray 集群已启动并包含所有节点

# 使用异步通信
with torch.cuda.stream(comm_stream):
    all_reduce(tensor)

# 同时在计算流上进行其他操作
with torch.cuda.stream(compute_stream):
    other_computation()

# 对于小张量，使用自定义 AllReduce
# vLLM 会自动选择最优策略

# 手动指定层分配（如果需要）
# 默认情况下 vLLM 会均匀分配

from vllm.distributed import (
    get_tensor_model_parallel_rank,
    get_tensor_model_parallel_world_size,
    get_pipeline_model_parallel_rank,
    get_pipeline_model_parallel_world_size,
)

# 打印当前进程的并行信息
print(f"TP Rank: {get_tensor_model_parallel_rank()}")
print(f"TP World Size: {get_tensor_model_parallel_world_size()}")
print(f"PP Rank: {get_pipeline_model_parallel_rank()}")
print(f"PP World Size: {get_pipeline_model_parallel_world_size()}")

# 设置 NCCL 调试级别
export NCCL_DEBUG=INFO

# 设置 NCCL 超时
export NCCL_TIMEOUT=1800

# 禁用 P2P 通信（调试用）
export NCCL_P2P_DISABLE=1