# vllm/v1/spec_decode/llm_base_proposer.py
class SpecDecodeBaseProposer:
    def propose(self, ...):
        # 第一步 forward
        self.set_inputs_first_pass(...)         # 准备第一步的 input_ids / positions
        self.model(...)                         # draft 模型第一步 forward
        self._sample_draft_tokens(...)          # 采样第一个 draft token

        # 剩余 K-1 步循环
        for token_index in range(K - 1):
            self._update_positions_dependent_metadata(...)  # positions / slot_mapping + 1
            self.model(...)
            self._sample_draft_tokens(...)

        return torch.stack(draft_tokens)        # [batch_size, K]

# vllm/v1/spec_decode/medusa.py
class MedusaProposer:
    def propose(self, ...):
        blocks = self.model(target_hidden_states)
        logits = self.model.compute_logits(blocks)
        return argmax(logits)  # [batch_size, num_heads]

# vllm/v1/worker/gpu/spec_decode/speculator.py
class BaseSpeculator(ABC):
    def init_cudagraph_manager(self, cudagraph_mode): ...
    def capture(self, attn_states): ...
    def propose(self,
                last_hidden_states, aux_hidden_states,
                num_sampled, num_rejected,
                last_sampled, next_prefill_tokens,
                temperature, seeds, ...): ...

# vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
class AutoRegressiveSpeculator(DraftModelSpeculator):
    def propose(self, ...):
        self.prepare_prefill_inputs(...)    # Triton kernel：shift input_ids，定位 last_token_indices
        self._prefill(...)                  # draft 模型第 0 步 forward，写 KV，生成 draft_tokens[:,0]
        self.prepare_decode_inputs(...)     # positions/seq_lens + 1，input_ids ← draft_tokens[:,0]
        self._multi_step_decode(...)

    def _multi_step_decode(self, ...):
        for step in range(1, K):
            self._build_draft_attn_metadata(...)  # 重算 slot mapping
            self._run_model(...)                   # draft forward，写新 KV
            self.sample_draft(...)                 # argmax / gumbel
            self.update_draft_inputs(...)          # Triton kernel：更新 input_ids/positions

# vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
# _rejection_kernel（Triton kernel，第 459 行，每个请求一个 warp）

@triton.jit
def _rejection_kernel(...):
    for i in range(num_draft_tokens):
        # 标准 rejection sampling（非 greedy，Leviathan et al. 2023）
        accepted &= target_log_prob > log(u) + draft_log_prob  # 概率比检验

        # greedy 模式
        target_argmax = compute_global_target_argmax(...)
        accepted &= (target_argmax == draft_sampled)            # 精确匹配

    rejected_steps[req] = first_rejected_index

# my_proposer.py
from vllm.v1.spec_decode.ngram_proposer import NgramProposer

class MyPrefixProposer:
    """
    在 prompt token 里找与 context 末尾匹配的前缀，
    取后续 num_speculative_tokens 个 token 作为 draft。
    """
    def __init__(self, vllm_config, device):
        self.num_speculative_tokens = (
            vllm_config.speculative_config.num_speculative_tokens
        )

    def load_model(self, target_model, target_attn_layer_names):
        pass  # 不需要模型

    def propose(
        self,
        requests,              # List[CachedRequestState]
        token_ids_cpu,         # numpy array，所有请求的 token 序列
        position_ids_cpu,
        input_batch,
        ...
    ):
        draft_token_ids = []
        num_draft_tokens = []

        for req in requests:
            ctx = token_ids_cpu[req.start : req.end]  # 当前 context
            prompt = ctx[: req.prompt_len]            # prompt 部分
            suffix = ctx[-3:]                         # 取末尾 3 token 做匹配 key

            # 在 prompt 里找 suffix，取后续 token 作为 draft
            drafts = self._find_continuation(prompt, suffix,
                                             self.num_speculative_tokens)
            draft_token_ids.extend(drafts)
            num_draft_tokens.append(len(drafts))

        return SpecDecodeMetadata(
            draft_token_ids=torch.tensor(draft_token_ids, device="cuda"),
            num_draft_tokens=num_draft_tokens,
            ...
        )

    def _find_continuation(self, prompt, suffix, k):
        for i in range(len(prompt) - len(suffix)):
            if (prompt[i : i + len(suffix)] == suffix).all():
                return prompt[i + len(suffix) : i + len(suffix) + k].tolist()
        return []

# vllm/v1/spec_decode/draft_model.py（参考）
class DraftModelProposer(SpecDecodeBaseProposer):

    @override
    def _create_draft_vllm_config(self) -> VllmConfig:
        # 把 draft 模型的 model_config / parallel_config 替换进去
        return replace(base, model_config=spec.draft_model_config, ...)

    @override
    def _get_model(self) -> nn.Module:
        # 加载 draft 模型权重
        return get_model(vllm_config=self._create_draft_vllm_config(), ...)

    @override
    def _maybe_share_lm_head(self, target_language_model: nn.Module) -> None:
        pass  # DraftModel 不共享权重；EAGLE 子类这里会做权重绑定

# serve/attention/backend.py
class BatchTopKAttention:
    # Draft 模式：REGISTER_INPUT_TRANSFORM 把 KV index 重定向到 flatten_indices（稀疏索引）
    core_variant_decl = """
    REGISTER_INPUT_TRANSFORM(sparse_kv_transform, {
        if (request_type[request_idx] == DRAFT) {
            kv_idx = flatten_indices[kv_idx];  // 只读 Top-K token 的 KV
        }
    });
    """

class BatchAttentionScore:
    # Verify 模式：REGISTER_LOGITS_TRANSFORM 把注意力 logits dump 出来
    score_variant_decl = """
    REGISTER_LOGITS_TRANSFORM(dump_attn_scores, {
        if (request_type[request_idx] == VERIFY) {
            dump_logits[...] = logits;  // 供后续 Top-K 更新使用
        }
    });
    """

# serve/request/kv_cache_ptr/pillar.py
class PillarCachePtr(StreamingCachePtr):

    def dispatch_prepare_fwd_metadata(self, ...):
        # 每步 draft 前：把 Top-K 索引打包成 flatten_indices，传给注意力 kernel
        self.concat_kv_indices_contiguous_kernel(
            self.selected_indices_buffer_gpu,  # Top-K token 索引
            self.recent_indices,               # 最近 token 索引（anchor）
            self.flatten_indices,              # 输出：拼接后的稀疏索引
        )

    def dispatch_update_selected_indices(self, ...):
        # 每次 verify 后：用 dump_logits 跑 Top-K，更新 Pillars
        self.topk_wrapper.run(
            self.dump_logits_snapshot,         # verify 阶段 dump 的注意力分数
            self.selected_indices_buffer_gpu,  # 更新 Top-K 索引
            k=self.num_selected_tokens,
        )

# serve/request/request.py
class ReqExecType(Enum):
    NORMAL = 0   # 普通 prefill
    DRAFT  = 1   # 稀疏 KV draft
    VERIFY = 2   # 全量 KV verify + rejection sampling
    STALL  = 3   # 等待 CPU 侧 Top-K 计算完成