源码精读：ModelRunner —— GPU 前向引擎

Scheduler.run_batch
   └── self.model_worker.forward_batch_generation(batch)
          └── ModelRunner.forward(forward_batch)        ← 本章主角
                ├── forward_extend(...)
                ├── forward_decode(...)
                └── forward_idle(...)

class ForwardBatch:
    forward_mode: ForwardMode                       # EXTEND / DECODE / IDLE
    batch_size: int
    input_ids: torch.Tensor                         # [total_tokens]
    req_pool_indices: torch.Tensor                  # [batch_size]
    seq_lens: torch.Tensor                          # [batch_size]
    out_cache_loc: torch.Tensor                     # 新算 KV 写到哪些 slot
    token_to_kv_pool: TokenToKVPool                 # KV cache 容器
    req_to_token_pool: ReqToTokenPool               # 逻辑 token → 物理 slot 映射
    attn_backend: AttentionBackend                  # 选哪个 backend
    sampling_info: SamplingBatchInfo                # 采样参数（一 batch 多种）
    # ... 多模态、LoRA、speculative 等可选字段

class ForwardMode(IntEnum):
    PREFILL = 1
    EXTEND  = 2      # 等同 prefill（变种）
    DECODE  = 3
    IDLE    = 4
    ...

def __init__(self, server_args, ...):
    self.load_model()                      # 1. 加载权重（1236）
    self.init_attention_backend()          # 2. 选 attention backend
    self.init_lora()                       # 3. LoRA
    self.init_token_to_kv_pool()           # 4. KV cache 容器
    self.init_req_to_token_pool()          # 5. 请求 → token 映射
    self.init_sampler()                    # 6. 采样器
    if cuda_graph_enabled:
        self.init_cuda_graphs()             # 7. 录 CUDA Graph

def load_model(self):
    # 1) 根据 model_config 找对应模型类
    model_cls = get_model_architecture(self.model_config)  # 在 srt/models/ 里

    # 2) 实例化（先在 meta device 上，省显存）
    with init_empty_weights():
        self.model = model_cls(config=self.model_config, ...)

    # 3) 把权重逐层 load 到 GPU
    if self.server_args.load_format == "auto":
        load_safetensors_weights(self.model, self.model_path)
    else:
        ...

    # 4) post-process: tie weights, init RoPE, ...
    self.model.post_init()

def forward(self, forward_batch: ForwardBatch, ...):
    self.attn_backend.init_forward_metadata(forward_batch)   # 预算 attention 元数据

    if forward_batch.forward_mode == ForwardMode.DECODE:
        return self.forward_decode(forward_batch)
    elif forward_batch.forward_mode == ForwardMode.IDLE:
        return self.forward_idle(forward_batch)
    elif forward_batch.forward_mode.is_extend():
        return self.forward_extend(forward_batch)
    else:
        raise ValueError(...)

def forward_extend(self, forward_batch):
    # 不走 CUDA Graph（prompt 长度可变）
    with torch.no_grad():
        hidden_states = self.model.forward(
            forward_batch.input_ids,
            forward_batch.positions,
            forward_batch,             # attention backend 用
        )
        # 取每个请求最后一个 token 的 hidden state，算 logits
        last_token_hs = self.extract_last_token_hidden(hidden_states, forward_batch)
        logits = self.model.compute_logits(last_token_hs)
        next_tokens = self.sampler.sample(logits, forward_batch.sampling_info)

    return ModelRunnerOutput(
        next_token_ids=next_tokens,
        # ... logprobs, embeddings 等可选
    )

def forward_decode(self, forward_batch):
    if cuda_graph_can_replay(forward_batch):
        return self.cuda_graph_runner.replay(forward_batch)   # ★ 走 graph
    else:
        # 回退到非 graph 路径（batch size 没匹配的 graph）
        with torch.no_grad():
            hidden = self.model.forward(...)
            logits = self.model.compute_logits(hidden)
            next_tokens = self.sampler.sample(logits, ...)
        return ModelRunnerOutput(next_token_ids=next_tokens)

def forward_idle(self, forward_batch):
    # 没有真请求，但要让 CUDA stream 不空转（保持 metric 节奏）
    ...
    return ModelRunnerOutput(...)

def init_cuda_graphs(self):
    # 给一组常见 batch size 都录一份 graph
    capture_batch_sizes = [1, 2, 4, 8, 16, 32, 64, 128, 256]   # 可配
    for bs in capture_batch_sizes:
        dummy_batch = make_dummy_decode_batch(bs)
        with torch.cuda.graph(graph) as g:
            self.model.forward(dummy_batch.input_ids, ...)
        self.cuda_graph_pool[bs] = g

def forward_batch_generation(self, batch, on_publish=None):
    with self.forward_stream_ctx:
        # 在 forward_stream 上跑 model
        output = self.forward(batch)
    if on_publish:
        on_publish(...)                                # spec_v2 模式中途回调
    return output

def update_weights_from_disk(self, recv_req):
    ...

def update_weights_from_distributed(self, recv_req):
    ...

def update_weights_from_tensor(self, recv_req):
    ...

def update_weights_from_ipc(self, recv_req):
    ...