xLLM NPU 推理执行主干

0. 一页全景

从 scheduler 攒出一个 batch，到把采样结果写回 sequence，中间是一条清晰的四段流水线。先建立这张整体图，后面各节再逐段展开。

Scheduler 攒出 Batch（一组 Sequence）
        │
        ▼  Executor::prepare_inputs(batch)
┌───────────────────────────────────────────────────────────┐
│ prepare 阶段（host → device 元数据构造）                     │
│  Batch::prepare_forward_input                               │
│   └ BatchInputBuilder::build_forward_input                  │
│       ├ process_sequences   逐序列攒 host vector（可多线程） │
│       │   ├ extract_tokens_and_positions  token/position    │
│       │   └ setup_kv_cache_info   slot_mapping/block_tables  │
│       ├ padding_decode_batch_size                           │
│       └ state_to_forward_input   host vector → ForwardInput  │
└───────────────────────────────────────────────────────────┘
        │  ForwardInput（CPU host tensors）
        ▼  WorkerImpl::step → step_internal
┌───────────────────────────────────────────────────────────┐
│ prepare_work_before_execute（在 prepare_stream 上）          │
│  ├ to(device)   H2D，host → NPU                             │
│  ├ CP/DP/EP padding                                         │
│  ├ apply_kv_block_swaps                                     │
│  └ 各类模型专有 device 侧准备                                │
│  → record_event()  给 forward 流做跨流 barrier              │
└───────────────────────────────────────────────────────────┘
        │  ForwardInput（device tensors）
        ▼  Executor::forward → ExecutorImpl::run
┌───────────────────────────────────────────────────────────┐
│ forward 阶段                                                │
│  base: model_->forward(...)  eager                          │
│  acl_graph: decode / spec-verify → capture-or-replay        │
│  → ModelOutput{ hidden_states, ... }                        │
└───────────────────────────────────────────────────────────┘
        │  hidden_states
        ▼
   model_->logits(hidden, selected_token_idxes)  只取要采样的行
        │  logits
        ▼  Sampler::forward
   penalties → top-k/top-p → softmax → greedy/random/mixed
        │  SampleOutput{ next_tokens, logprobs, ... }
        ▼
   ForwardOutput  回 scheduler（写回 token、推进 sequence）

1. 入口：step 链

WorkerImpl::step 是纯虚接口，LLM 的实体在 LLMWorkerImpl::step：取当前 stream → 等 input ready events → 进入 step_internal。step_internal 是真正的主干：

1. （PD PUSH 模式）先起一个 kv_cache 传输 future。
2. forward：model_executor_->forward(token_ids, positions, kv_caches_, input_params) → ModelOutput。hidden_states 未定义则返回 nullopt。
3. logits：若 sampling_params.selected_token_idxes 定义，只对需要采样的行算 lm_head，省算力。
4. sample：sampler_->forward(logits, sampling_params)，可选 beam search。
5. speculative：开启投机解码时走 MTP 分支（见 §6）。
6. 非 driver / 非 overlap / 非 spec 的纯计算 worker 提前返回 nullopt，只有 driver 组装 ForwardOutput。

2. prepare 阶段（一）：host 张量构造

入口委托链：Executor::prepare_inputs → BaseExecutorImpl::prepare_inputs → Batch::prepare_forward_input → BatchInputBuilder::build_forward_input。ACL Graph 的 executor 覆写了 prepare_inputs 但实现与 base 一致——分桶逻辑在 run() 里，不在 prepare。

2.1 process_sequences

逐 Sequence 调 process_single_sequence；序列数达到线程池规模时多线程并行，每线程独立 BuilderState，最后按 offset 合并。核心两个子步：

extract_tokens_and_positions —— prefill / decode 分叉的根：

for (uint32_t j = n_kv_cache_tokens; j < seq_len; ++j) {
  state.flatten_tokens_vec.emplace_back(token_ids[j]);
  state.flatten_positions_vec.push_back(j);
}

• prefill：n_kv_cache_tokens = 0（或 = 已命中前缀长度），取 [n_kv_cache_tokens, seq_len) 整段 prompt。
• decode：n_kv_cache_tokens = seq_len - 1，只取最后 1 个 token。
• 采样选择器：非采样请求只在最后一个 prompt token 落一个 sampling entry；采样请求按 sample_slots 落多个。
• CP 右 padding、MTP bootstrap embedding 也在这里。

setup_kv_cache_info —— KV 寻址元数据：

• incr_kv_cache_tokens_num(q_seq_len) 推进该序列已缓存 token 数。
• kv_cache_slots(...) → new_token_slot_ids，即 slot_mapping：新 token 写进物理 cache 的槽位。
• 遍历 blocks 攒 block_ids → block_tables_vec，并填 paged_kv 三件套（indptr / indices / last_page_len）。

2.2 padding_decode_batch_size

decode 批按 min_decoding_batch_size 补齐，对 graph 分桶友好。

2.3 state_to_forward_input

把所有 host vector 装成 ForwardInput（此时还是 CPU tensor），含 token_ids、positions、input_params（attention metadata、parallel、embedding 等子结构）、sampling_params。

3. prepare 阶段（二）：H2D + 设备侧准备

WorkerImpl::prepare_work_before_execute 全程在 prepare_stream_ 上，核心是内部 lambda prepare_device_on_stream：

1. processed_input = prep_for_device.to(device_, dtype_)：整个 ForwardInput 一次性 H2D。
2. ensure_forward_input_device_tensors：补齐 device 张量。
3. CP prefill 侧（cp_size > 1 且非 decode）：prefill 输入准备、CP/EP padding、KV split 重映射（非幂等，靠 !cp_partitioned 保证只跑一次）。
4. 空 shard 造假输入，避免多卡某 rank 空批挂住集合通信。
5. apply_kv_block_swaps：换页 / 拷贝。
6. MLA：chunked prefill 下的 prefix-cache 输入准备。
7. DP/EP padding。
8. 模型专有 device 侧准备：例如线性注意力（GDN）类模型在此恢复 recurrent state slot。

4. forward 阶段：eager vs ACL Graph

Executor::forward → impl_->run(...)。backend 选择：开启 graph 且非 vlm 时，NPU 下走 AclGraphExecutorImpl；否则走 base（eager）/ vlm。

4.1 base（eager）

直接 model_->forward(tokens, positions, kv_caches, params)。

4.2 ACL Graph

何时用 graph（否则 eager fallback）：decode 阶段，或 spec-verify + chunked prefill 的特定模型。fallback 条件包括：普通 chunked prefill、单层模型、DP 多 rank 非全 decode、batch 超过 decode 上限、序列超过最大位置。

bucket 分桶：按 num_tokens 归入 1 / 2 / 4 / 8 / 向上取整到 16 的倍数；关闭 padding 时每个真实 token 数独立建图。

capture（首次遇到某 bucket key）:
  1. 把 tokens/positions/block_tables 等拷进预分配的持久 buffer，
     返回引用这些固定地址的 graph_params
  2. 切到 high-priority capture_stream → capture_begin
     → model->forward(persistent_tokens, ..., graph_params)
     → set_hidden_states → capture_end
  3. test replay + 让当前流等待 graph

replay（命中 bucket）:
  1. 把当前 batch 的实际数据 non_blocking 写入
     capture 时记录的同一持久 buffer 地址 ← 绑定机制
  2. graph.replay()  复用 mempool 内全部临时 tensor，重跑算子序列
  3. 让当前流等待 graph → 返回 hidden_states 切片

5. sample 阶段

Sampler::forward 的处理顺序：

1. frequency / presence penalty、repetition penalty。
2. 若 selected_token_idxes != sample_idxes，index_select 出要采样的子集。
3. 可选 filter_mask 相加。
4. apply_top_k_top_p。
5. softmax → 按 all-random / all-greedy / mixed（torch::where(do_sample, random, greedy)）出 samples。
6. logprobs / top_logprobs。

输出 SampleOutput{ next_tokens, probs, logprobs, top_tokens, top_logprobs }，装进 ForwardOutput 回 scheduler。

6. MTP 投机分支

结构上是 MTPWorkerImpl → SpeculativeWorkerImpl → WorkerImpl 的继承链。SpeculativeWorkerImpl 持 target worker；MTPWorkerImpl 额外持 draft worker（1 层）+ embedding_cache_。入口 SpeculativeWorkerImpl::step 按 empty / prefill / decode 分发。

draft 阶段（循环 num_speculative_tokens 次）

• 第 1 步：从 EmbeddingCache 取上一 step 的 target hidden state 作 embedding，构造双行 extend 输入。
• 第 k>0 步：把上一 draft 输出 next_tokens 喂回，hidden 也取自上一 draft 输出。
• draft KV cache 独立按固定容量分配，不走 prefix cache 查找；draft probs 每步压成 selected-only。

verify 阶段

• 每序列展开 num_speculative_tokens + 1 行，token = [current, -1, -2, ...]（负数 placeholder）。
• placeholder 被实际 draft token 覆盖后，target 一次 forward 并行 verify。

接受 / 拒绝

• 随机路径：acceptance_prob = target_p / draft_p，逐位与 uniform 比；拒绝则从 recovered_probs = max(0, target_p - draft_p) 归一重采。greedy 直接比 argmax。
• bonus token = target 对最后一个 validate position 的采样，无条件附在接受串末。
• accepted tokens + embeddings 写回 EmbeddingCache，供下一 step draft 用。KV slot 级回滚由 scheduler 按实际 accepted 数处理，worker 代码无显式 rollback。

7. 流（stream）模型小结

step 开头 wait_input_ready_events 等 input 的 metadata_ready_event。开启 graph 时 prepare 全程持 DeviceCaptureLock，防 capture 被异步 H2D 打断。