nano-vllm 源码阅读

nano-vllm
├── bench.py
├── example.py
├── LICENSE
├── nanovllm
│   ├── config.py
│   ├── engine
│   │   ├── block_manager.py
│   │   ├── llm_engine.py
│   │   ├── model_runner.py
│   │   ├── scheduler.py
│   │   └── sequence.py
│   ├── __init__.py
│   ├── layers
│   │   ├── activation.py
│   │   ├── attention.py
│   │   ├── embed_head.py
│   │   ├── layernorm.py
│   │   ├── linear.py
│   │   ├── rotary_embedding.py
│   │   └── sampler.py
│   ├── llm.py
│   ├── models
│   │   └── qwen3.py
│   ├── sampling_params.py
│   └── utils
│       ├── context.py
│       └── loader.py
├── pyproject.toml
└── README.md

tokenizer = AutoTokenizer.from_pretrained(path)
llm = LLM(path, enforce_eager=True, tensor_parallel_size=1) # LLM 就是 LLMEngine 的一层名字上的包装

prompts = [
    "introduce yourself",
    "list all prime numbers within 100",
]
prompts = [
    tokenizer.apply_chat_template(
        [{"role": "user", "content": prompt}],
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=True
    )
    for prompt in prompts
]

# nanovllm/sampling_params.py
from dataclasses import dataclass

@dataclass
class SamplingParams:
    temperature: float = 1.0
    max_tokens: int = 64
    ignore_eos: bool = False

# nanovllm/engine/scheduler.py, line 68:
if (not seq.ignore_eos and token_id == self.eos) or seq.num_completion_tokens == seq.max_tokens:
    seq.status = SequenceStatus.FINISHED
    self.block_manager.deallocate(seq)
    self.running.remove(seq)

# nanovllm/engine/sequence.py, line 42:
@property
def num_completion_tokens(self):
    return self.num_tokens - self.num_prompt_tokens

# nanovllm/engine/scheduler.py, line 68
if (not seq.ignore_eos and token_id == self.eos) or seq.num_completion_tokens == seq.max_tokens:
    seq.status = SequenceStatus.FINISHED
    self.block_manager.deallocate(seq)
    self.running.remove(seq)

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta")
chat = [
  {"role": "user", "content": "Hello, how are you?"},
  {"role": "assistant", "content": "I'm doing great. How can I help you today?"},
  {"role": "user", "content": "I'd like to show off how chat templating works!"},
]

print(tokenizer.apply_chat_template(chat, tokenize=False))

<|user|>
Hello, how are you?</s>
<|assistant|>
I'm doing great. How can I help you today?</s>
<|user|>
I'd like to show off how chat templating works!</s>

[523, 28766, 1838, 28766, 28767, 13, 16230, 28725, 910, 460, 368, 28804, 2, 28705, 13, 28789, 28766, 489, 11143, 28766, 28767, 13, 28737, 28742, 28719, 2548, 1598, 28723, 1602, 541, 315, 1316, 368, 3154, 28804, 2, 28705, 13, 28789, 28766, 1838, 28766, 28767, 13, 28737, 28742, 28715, 737, 298, 1347, 805, 910, 10706, 5752, 1077, 3791, 28808, 2, 28705, 13]

print(tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=False))
# 输出：
# <|user|>
# Hello, how are you?</s>
# <|assistant|>
# I'm doing great. How can I help you today?</s>
# <|user|>
# I'd like to show off how chat templating works!</s>

print(tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True))
# 输出：
# <|user|>
# Hello, how are you?</s>
# <|assistant|>
# I'm doing great. How can I help you today?</s>
# <|user|>
# I'd like to show off how chat templating works!</s>
# <|assistant|>

# nanovllm/llm.py
from nanovllm.engine.llm_engine import LLMEngine

class LLM(LLMEngine):
    pass

import os
from dataclasses import dataclass
from transformers import AutoConfig

@dataclass
class Config:
    model: str
    max_num_batched_tokens: int = 16384
    max_num_seqs: int = 512
    max_model_len: int = 4096
    gpu_memory_utilization: float = 0.9
    tensor_parallel_size: int = 1
    enforce_eager: bool = False
    hf_config: AutoConfig | None = None
    eos: int = -1
    kvcache_block_size: int = 256
    num_kvcache_blocks: int = -1

    def __post_init__(self):
        assert os.path.isdir(self.model)
        assert self.kvcache_block_size % 256 == 0
        assert 1 <= self.tensor_parallel_size <= 8
        self.hf_config = AutoConfig.from_pretrained(self.model)
        self.max_model_len = min(self.max_model_len, self.hf_config.max_position_embeddings)
        assert self.max_num_batched_tokens >= self.max_model_len

self.ps = [] # 所有子进程对象，方便管理（终止、join 等）
self.events = [] # 所有事件对象，用于进程间同步
ctx = mp.get_context("spawn") # 得到多进程上下文，使用 spawn 模式来避免 fork 时的 CUDA 问题
for i in range(1, config.tensor_parallel_size): # 注意是从 1 开始，因为当前进程（父进程）0 占用了一个张量并行的位置
    event = ctx.Event()
    process = ctx.Process(target=ModelRunner, args=(config, i, event)) # 创建子进程，将会运行 ModelRunner(config, i, event)
    process.start()
    self.ps.append(process)
    self.events.append(event) # 收集子进程和事件对象

self.model_runner = ModelRunner(config, 0, self.events) # 注意这里需要初始化当前进程的 model_runner，最后一个参数传入的是 self.events 列表而非单个事件，这是因为 rank 0 负责协调，能访问所有事件
self.tokenizer = AutoTokenizer.from_pretrained(config.model, use_fast=True)
config.eos = self.tokenizer.eos_token_id
self.scheduler = Scheduler(config)

atexit.register(self.exit)

def exit(self):
    self.model_runner.call("exit")
    del self.model_runner
    for p in self.ps:
        p.join()

def add_request(self, prompt: str | list[int], sampling_params: SamplingParams):
    if isinstance(prompt, str):
        prompt = self.tokenizer.encode(prompt)
    seq = Sequence(prompt, sampling_params)
    self.scheduler.add(seq)

def __init__(self, token_ids: list[int], sampling_params = SamplingParams()):
    self.seq_id = next(Sequence.counter)
    self.status = SequenceStatus.WAITING
    self.token_ids = copy(token_ids)
    self.last_token = token_ids[-1]
    self.num_tokens = len(self.token_ids)
    self.num_prompt_tokens = len(token_ids)
    self.num_cached_tokens = 0
    self.block_table = []
    self.temperature = sampling_params.temperature
    self.max_tokens = sampling_params.max_tokens
    self.ignore_eos = sampling_params.ignore_eos

def __getstate__(self):
    return (self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table,
            self.token_ids if self.num_completion_tokens == 0 else self.last_token)

def __setstate__(self, state):
    self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table = state[:-1]
    if self.num_completion_tokens == 0:
        self.token_ids = state[-1]
    else:
        self.last_token = state[-1]

# read_shm
method_name, *args = pickle.loads(self.shm.buf[4:n+4])
# write_shm
data = pickle.dumps([method_name, *args])
n = len(data)
self.shm.buf[0:4] = n.to_bytes(4, "little")
self.shm.buf[4:n+4] = data

# nanovllm/engine/llm_engine.py, line 59
def generate(
    self,
    prompts: list[str] | list[list[int]],
    sampling_params: SamplingParams | list[SamplingParams],
    use_tqdm: bool = True,
) -> list[str]:
    if use_tqdm:
        pbar = tqdm(total=len(prompts), desc="Generating", dynamic_ncols=True)
    if not isinstance(sampling_params, list):
        sampling_params = [sampling_params] * len(prompts) # 这里处理不是列表的情况，默认所有序列生成参数相同
    for prompt, sp in zip(prompts, sampling_params):
        self.add_request(prompt, sp) # 调用上一部分提到的 add_request 方法
    outputs = {}
    prefill_throughput = decode_throughput = 0.
    while not self.is_finished():
        t = perf_counter()
        output, num_tokens = self.step() # step 方法，执行一步推理，我们在下一部分来看
        if use_tqdm:
            if num_tokens > 0: # 约定大于 0 时是 prefill，小于 0 时是 decode
                prefill_throughput = num_tokens / (perf_counter() - t)
            else:
                decode_throughput = -num_tokens / (perf_counter() - t)
            pbar.set_postfix({
                "Prefill": f"{int(prefill_throughput)}tok/s",
                "Decode": f"{int(decode_throughput)}tok/s",
            })
        # 下面开始处理返回值，由于这些值和 step 方法相关，建议先看 step 的实现
        for seq_id, token_ids in output:
            outputs[seq_id] = token_ids
            if use_tqdm:
                pbar.update(1)
    outputs = [outputs[seq_id] for seq_id in sorted(outputs)]
    outputs = [{"text": self.tokenizer.decode(token_ids), "token_ids": token_ids} for token_ids in outputs]
    if use_tqdm:
        pbar.close()
    return outputs

# nanovllm/engine/llm_engine.py, line 48
def step(self):
    seqs, is_prefill = self.scheduler.schedule() # scheduler 负责规划每一步需要向前推理的请求，同时返回一个是否处于 prefill 阶段的 bool 值
    token_ids = self.model_runner.call("run", seqs, is_prefill) # 让实际模型跑一步得到新的 token ids
    self.scheduler.postprocess(seqs, token_ids) # 后处理
    outputs = [(seq.seq_id, seq.completion_token_ids) for seq in seqs if seq.is_finished] # 只返回已经完成的请求
    num_tokens = sum(len(seq) for seq in seqs) if is_prefill else -len(seqs) # 这里如果是 decode 阶段则会返回一个负数，和上面的 generate 方法中对应的 if-else 判断语句逻辑相符
    return outputs, num_tokens

# nanovllm/engine/llm_engine.py

# __init__
for i in range(1, config.tensor_parallel_size):
    event = ctx.Event()
    process = ctx.Process(target=ModelRunner, args=(config, i, event))
    process.start()
    self.ps.append(process)
    self.events.append(event)
self.model_runner = ModelRunner(config, 0, self.events)

# exit
self.model_runner.call("exit")
del self.model_runner

# step
token_ids = self.model_runner.call("run", seqs, is_prefill)

# nanovllm/engine/model_runner.py
class ModelRunner:
    def __init__(self, config: Config, rank: int, event: Event | list[Event]):
        ...
    def exit(self):
        ...
    def loop(self):
        ...
    def read_shm(self):
        ...
    def write_shm(self, method_name, *args):
        ...
    def call(self, method_name, *args):
        ...
    def warmup_model(self):
        ...
    def allocate_kv_cache(self):
        ...
    def prepare_block_tables(self, seqs: list[Sequence]):
        ...
    def prepare_prefill(self, seqs: list[Sequence]):
        ...
    def prepare_decode(self, seqs: list[Sequence]):
        ...
    def prepare_sample(self, seqs: list[Sequence]):
        ...
    def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill: bool):
        ...
    def run(self, seqs: list[Sequence], is_prefill: bool) -> list[int]:
        ...
    def capture_cudagraph(self):
        ...

def __init__(self, config: Config, rank: int, event: Event | list[Event]):
    # 先设置 config 值
    self.config = config
    hf_config = config.hf_config
    self.block_size = config.kvcache_block_size
    self.enforce_eager = config.enforce_eager
    self.world_size = config.tensor_parallel_size
    self.rank = rank
    self.event = event

    # 启动进程组并绑定 CUDA 设备、设置 torch 数据类型
    dist.init_process_group("nccl", "tcp://localhost:2333", world_size=self.world_size, rank=rank)
    torch.cuda.set_device(rank)
    default_dtype = torch.get_default_dtype()
    torch.set_default_dtype(hf_config.torch_dtype)
    torch.set_default_device("cuda")
    # 加载模型和 sampler
    self.model = Qwen3ForCausalLM(hf_config)
    load_model(self.model, config.model)
    self.sampler = Sampler()
    # 预热，避免首次真实推理时出现不确定的延迟或内存不足的问题
    self.warmup_model()
    self.allocate_kv_cache()
    # CUDA graph 相关，暂时不考虑
    if not self.enforce_eager:
        self.capture_cudagraph()
    # 设置默认设备和 dtype
    torch.set_default_device("cpu")
    torch.set_default_dtype(default_dtype)

    # 如果采用了张量并行，那么就使用共享内存来实现不同进程之间的通讯，在 read_shm 和 write_shm 中有使用，并且 rank=0 的进程负责创建，其他进程直接连接
    if self.world_size > 1:
        if rank == 0:
            self.shm = SharedMemory(name="nanovllm", create=True, size=2**20)
            dist.barrier() # 先创建再等待其他进程同步
        else:
            dist.barrier() # 先同步得到 rank=0 进程创建的共享内存，然后再连接
            self.shm = SharedMemory(name="nanovllm")
            self.loop() # 其他进程进入 loop 开始循环，等待 rank=0 进程通过共享内存发送指令

def warmup_model(self):
    # 清空显存并重置峰值显存的记录
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    # 使用最大负载来得到实际运行时的显存峰值占用
    max_num_batched_tokens, max_model_len = self.config.max_num_batched_tokens, self.config.max_model_len
    num_seqs = min(max_num_batched_tokens // max_model_len, self.config.max_num_seqs)
    seqs = [Sequence([0] * max_model_len) for _ in range(num_seqs)] # 由于这次推理只需要起到预热的效果就行了，所以可以任意赋值
    self.run(seqs, True) # 跑一遍refill 阶段
    torch.cuda.empty_cache() # 清空显存

def allocate_kv_cache(self):
    config = self.config
    hf_config = config.hf_config
    free, total = torch.cuda.mem_get_info()
    used = total - free
    peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"]
    current = torch.cuda.memory_stats()["allocated_bytes.all.current"]
    num_kv_heads = hf_config.num_key_value_heads // self.world_size
    block_bytes = 2 * hf_config.num_hidden_layers * self.block_size * num_kv_heads * hf_config.head_dim * hf_config.torch_dtype.itemsize
    config.num_kvcache_blocks = int(total * config.gpu_memory_utilization - used - peak + current) // block_bytes
    assert config.num_kvcache_blocks > 0
    self.kv_cache = torch.zeros(2, hf_config.num_hidden_layers, config.num_kvcache_blocks, self.block_size, num_kv_heads, hf_config.head_dim)
    layer_id = 0
    for module in self.model.modules():
        if hasattr(module, "k_cache") and hasattr(module, "v_cache"):
            module.k_cache = self.kv_cache[0, layer_id]
            module.v_cache = self.kv_cache[1, layer_id]
            layer_id += 1

num_kv_heads = hf_config.num_key_value_heads // self.world_size
block_bytes = 2 * num_layers * block_size * num_kv_heads * head_dim * dtype_size

# config.num_kvcache_blocks = (可用显存字节) // block_bytes
config.num_kvcache_blocks = int(total * config.gpu_memory_utilization - used - peak + current) // block_bytes
assert config.num_kvcache_blocks > 0 # 断言必须 > 0，保证至少能分配一块。

[2, num_layers, num_kvcache_blocks, block_size, num_kv_heads, head_dim]

def loop(self):
    while True:
        method_name, args = self.read_shm()
        self.call(method_name, *args)
        if method_name == "exit":
            break

def read_shm(self):
    assert self.world_size > 1 and self.rank
    self.event.wait()
    n = int.from_bytes(self.shm.buf[0:4], "little")
    method_name, *args = pickle.loads(self.shm.buf[4:n+4])
    self.event.clear()
    return method_name, args

def write_shm(self, method_name, *args):
    assert self.world_size > 1 and not self.rank
    data = pickle.dumps([method_name, *args])
    n = len(data)
    self.shm.buf[0:4] = n.to_bytes(4, "little")
    self.shm.buf[4:n+4] = data
    for event in self.event:
        event.set()

def call(self, method_name, *args):
    if self.world_size > 1 and self.rank == 0:
        self.write_shm(method_name, *args)
    method = getattr(self, method_name, None)
    return method(*args)

@torch.inference_mode()
def capture_cudagraph(self):
    config = self.config
    hf_config = config.hf_config
    max_bs = min(self.config.max_num_seqs, 512)
    max_num_blocks = (config.max_model_len + self.block_size - 1) // self.block_size
    input_ids = torch.zeros(max_bs, dtype=torch.int64)
    positions = torch.zeros(max_bs, dtype=torch.int64)
    slot_mapping = torch.zeros(max_bs, dtype=torch.int32)
    context_lens = torch.zeros(max_bs, dtype=torch.int32)
    block_tables = torch.zeros(max_bs, max_num_blocks, dtype=torch.int32)
    outputs = torch.zeros(max_bs, hf_config.hidden_size)
    self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
    self.graphs = {}
    self.graph_pool = None

    for bs in reversed(self.graph_bs):
        graph = torch.cuda.CUDAGraph()
        set_context(False, slot_mapping=slot_mapping[:bs], context_lens=context_lens[:bs], block_tables=block_tables[:bs])
        outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # warmup
        with torch.cuda.graph(graph, self.graph_pool):
            outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # capture
        if self.graph_pool is None:
            self.graph_pool = graph.pool()
        self.graphs[bs] = graph
        torch.cuda.synchronize()
        reset_context()

    self.graph_vars = dict(
        input_ids=input_ids,
        positions=positions,
        slot_mapping=slot_mapping,
        context_lens=context_lens,
        block_tables=block_tables,
        outputs=outputs,
   )

for bs in reversed(self.graph_bs):  # 从大到小遍历 batch size，这样最大的内存需求会首先被满足
    graph = torch.cuda.CUDAGraph()
    set_context(...)  # 设置上下文信息
    outputs[:bs] = self.model(input_ids[:bs], positions[:bs])  # warmup，避免编译、初始化等一次性操作被记录，提高录制的 cuda graph 效率
    with torch.cuda.graph(graph, self.graph_pool): # 这个代码块中执行的所有CUDA操作都不会立即在GPU上运行，而是会被记录到 graph 对象中
        outputs[:bs] = self.model(input_ids[:bs], positions[:bs])  # capture
        if self.graph_pool is None:
            self.graph_pool = graph.pool() # 在成功捕获第一个 graph 后（即 bs 最大的 graph），保存其内存池，在后续继续捕获更小的 bs 的 graph 时共享，节约显存
        self.graphs[bs] = graph
        torch.cuda.synchronize() # 等待当前 graph 完全捕获
        reset_context()

def run(self, seqs: list[Sequence], is_prefill: bool) -> list[int]:
    input_ids, positions = self.prepare_prefill(seqs) if is_prefill else self.prepare_decode(seqs)
    temperatures = self.prepare_sample(seqs) if self.rank == 0 else None
    logits = self.run_model(input_ids, positions, is_prefill)
    token_ids = self.sampler(logits, temperatures).tolist() if self.rank == 0 else None
    reset_context()
    return token_ids

def prepare_prefill(self, seqs: list[Sequence]):
    input_ids = []
    positions = []
    cu_seqlens_q = [0]
    cu_seqlens_k = [0]
    max_seqlen_q = 0
    max_seqlen_k = 0
    slot_mapping = []
    block_tables = None
    for seq in seqs:
        seqlen = len(seq)
        input_ids.extend(seq[seq.num_cached_tokens:])
        positions.extend(list(range(seq.num_cached_tokens, seqlen)))
        seqlen_q = seqlen - seq.num_cached_tokens
        seqlen_k = seqlen
        cu_seqlens_q.append(cu_seqlens_q[-1] + seqlen_q)
        cu_seqlens_k.append(cu_seqlens_k[-1] + seqlen_k)
        max_seqlen_q = max(seqlen_q, max_seqlen_q)
        max_seqlen_k = max(seqlen_k, max_seqlen_k)
        if not seq.block_table:    # warmup
            continue
        for i in range(seq.num_cached_blocks, seq.num_blocks):
            start = seq.block_table[i] * self.block_size
            if i != seq.num_blocks - 1:
                end = start + self.block_size
            else:
                end = start + seq.last_block_num_tokens
            slot_mapping.extend(list(range(start, end)))
    if cu_seqlens_k[-1] > cu_seqlens_q[-1]:    # prefix cache
        block_tables = self.prepare_block_tables(seqs)
    input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
    positions = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
    cu_seqlens_q = torch.tensor(cu_seqlens_q, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
    cu_seqlens_k = torch.tensor(cu_seqlens_k, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
    slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
    set_context(True, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, slot_mapping, None, block_tables)
    return input_ids, positions

def prepare_decode(self, seqs: list[Sequence]):
    input_ids = []
    positions = []
    slot_mapping = []
    context_lens = []
    for seq in seqs:
        input_ids.append(seq.last_token)
        positions.append(len(seq) - 1)
        context_lens.append(len(seq))
        slot_mapping.append(seq.block_table[-1] * self.block_size + seq.last_block_num_tokens  - 1)
    input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
    positions = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
    slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
    context_lens = torch.tensor(context_lens, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
    block_tables = self.prepare_block_tables(seqs)
    set_context(False, slot_mapping=slot_mapping, context_lens=context_lens, block_tables=block_tables)
    return input_ids, positions

@dataclass
class Context:
    is_prefill: bool = False
    cu_seqlens_q: torch.Tensor | None = None
    cu_seqlens_k: torch.Tensor | None = None
    max_seqlen_q: int = 0
    max_seqlen_k: int = 0
    slot_mapping: torch.Tensor | None = None
    context_lens: torch.Tensor | None = None
    block_tables: torch.Tensor | None = None

for seq in seqs:
    seqlen = len(seq)
    input_ids.extend(seq[seq.num_cached_tokens:])
    positions.extend(list(range(seq.num_cached_tokens, seqlen)))
    seqlen_q = seqlen - seq.num_cached_tokens # 这里减掉已经缓存的 token 是因为只需要当前的 query 来计算 value，往前的 query 已经不需要了
    seqlen_k = seqlen
    cu_seqlens_q.append(cu_seqlens_q[-1] + seqlen_q) # 记录当前 seq 的终止位置，起始位置就是上一个 seq 的终止位置
    cu_seqlens_k.append(cu_seqlens_k[-1] + seqlen_k)
    max_seqlen_q = max(seqlen_q, max_seqlen_q) # 记录 query 最大长度
    max_seqlen_k = max(seqlen_k, max_seqlen_k)
    # block_table 为空则跳过，这是因为在 warmup_model 中还没有分配实际的缓存块表
    if not seq.block_table:    # warmup
        continue
    for i in range(seq.num_cached_blocks, seq.num_blocks): # 遍历未 cache 的物理块
        start = seq.block_table[i] * self.block_size
        if i != seq.num_blocks - 1: # 如果不是最后一个块
            end = start + self.block_size
        else:
            end = start + seq.last_block_num_tokens
        slot_mapping.extend(list(range(start, end))) # 添加计算出的物理 slot 索引范围

if cu_seqlens_k[-1] > cu_seqlens_q[-1]:    # prefix cache
    block_tables = self.prepare_block_tables(seqs)
# pin_memory=True 表示使用锁页内存，可以加速从 CPU 到 GPU 的数据传输
# non_blocking=True 表示可以不用等待数据传输完成就可以继续执行后续代码
input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
positions = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
cu_seqlens_q = torch.tensor(cu_seqlens_q, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
cu_seqlens_k = torch.tensor(cu_seqlens_k, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
set_context(True, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, slot_mapping, None, block_tables)

def prepare_block_tables(self, seqs: list[Sequence]):
    max_len = max(len(seq.block_table) for seq in seqs) # 最大 block 块数
    block_tables = [seq.block_table + [-1] * (max_len - len(seq.block_table)) for seq in seqs] # 将较短的 block table 右补齐至最长的长度
    block_tables = torch.tensor(block_tables, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) # 得到对齐后的 block table
    return block_tables

def set_context(is_prefill, cu_seqlens_q=None, cu_seqlens_k=None, max_seqlen_q=0, max_seqlen_k=0, slot_mapping=None, context_lens=None, block_tables=None):
    global _CONTEXT
    _CONTEXT = Context(is_prefill, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, slot_mapping, context_lens, block_tables)

def prepare_decode(self, seqs: list[Sequence]):
    input_ids = []
    positions = []
    slot_mapping = []
    context_lens = []
    for seq in seqs:
        input_ids.append(seq.last_token)
        positions.append(len(seq) - 1)
        context_lens.append(len(seq)) # 上下文总长度
        slot_mapping.append(seq.block_table[-1] * self.block_size + seq.last_block_num_tokens  - 1)
    input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
    positions = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
    slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
    context_lens = torch.tensor(context_lens, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
    block_tables = self.prepare_block_tables(seqs) # 此时必然已经分配了 kv blocks 了，所以一定会需要对齐
    set_context(False, slot_mapping=slot_mapping, context_lens=context_lens, block_tables=block_tables) # 由于 decode 阶段固定只生成一个 token，因此 cu_seqlens_q, max_seqlen_q 等变量不再需要
    return input_ids, positions

temperatures = self.prepare_sample(seqs) if self.rank == 0 else None

logits = self.run_model(input_ids, positions, is_prefill)

token_ids = self.sampler(logits, temperatures).tolist() if self.rank == 0 else None
reset_context()

@torch.inference_mode()
def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill: bool):
    if is_prefill or self.enforce_eager or input_ids.size(0) > 512:
        return self.model.compute_logits(self.model(input_ids, positions))
    else:
        bs = input_ids.size(0)
        context = get_context()
        graph = self.graphs[next(x for x in self.graph_bs if x >= bs)]
        graph_vars = self.graph_vars
        graph_vars["input_ids"][:bs] = input_ids
        graph_vars["positions"][:bs] = positions
        graph_vars["slot_mapping"].fill_(-1)
        graph_vars["slot_mapping"][:bs] = context.slot_mapping
        graph_vars["context_lens"].zero_()
        graph_vars["context_lens"][:bs] = context.context_lens
        graph_vars["block_tables"][:bs, :context.block_tables.size(1)] = context.block_tables
        graph.replay()
        return self.model.compute_logits(graph_vars["outputs"][:bs])

bs = input_ids.size(0)
context = get_context()
graph = self.graphs[next(x for x in self.graph_bs if x >= bs)] # 找到刚好超过 batch size 的预捕获值，下面的过程都是在“重现”当时捕获的场景
graph_vars = self.graph_vars
graph_vars["input_ids"][:bs] = input_ids
graph_vars["positions"][:bs] = positions
graph_vars["slot_mapping"].fill_(-1)
graph_vars["slot_mapping"][:bs] = context.slot_mapping
graph_vars["context_lens"].zero_()
graph_vars["context_lens"][:bs] = context.context_lens
graph_vars["block_tables"][:bs, :context.block_tables.size(1)] = context.block_tables
graph.replay() # 触发 GPU 执行之前被捕捉和编译好的完整计算图，相当于用一个 CPU 指令就启动了模型中所有层、所有 CUDA Kernel 的执行
return self.model.compute_logits(graph_vars["outputs"][:bs])

# __init__
self.scheduler = Scheduler(config)

# add_request
self.scheduler.add(seq)

# step
seqs, is_prefill = self.scheduler.schedule()
...
self.scheduler.postprocess(seqs, token_ids)

# is_finished
return self.scheduler.is_finished()

def __init__(self, config: Config): self.max_num_seqs = config.max_num_seqs
    self.max_num_batched_tokens = config.max_num_batched_tokens
    self.eos = config.eos
    self.block_manager = BlockManager(config.num_kvcache_blocks, config.kvcache_block_size)
    self.waiting: deque[Sequence] = deque() # 等待被执行的序列
    self.running: deque[Sequence] = deque() # 正在推理的序列

def is_finished(self):
    return not self.waiting and not self.running

def add(self, seq: Sequence):
    self.waiting.append(seq)

def schedule(self) -> tuple[list[Sequence], bool]:
    # prefill
    scheduled_seqs = []
    num_seqs = 0
    num_batched_tokens = 0
    while self.waiting and num_seqs < self.max_num_seqs:
        seq = self.waiting[0] # 得到 waiting 队列中第一个，但暂时没有移出 waiting 队列
        if num_batched_tokens + len(seq) > self.max_num_batched_tokens or not self.block_manager.can_allocate(seq): # 如果超出最大限度，直接停止 schedule
            break
        num_seqs += 1
        self.block_manager.allocate(seq) # 分配新的 kv cache 块
        num_batched_tokens += len(seq) - seq.num_cached_tokens
        seq.status = SequenceStatus.RUNNING
        self.waiting.popleft() # 此时才将该请求移出 waiting 队列
        self.running.append(seq)
        scheduled_seqs.append(seq)
    # 由于上面进入 scheduled_seqs 的请求只能是 waiting 序列中的前若干项，所以必定是 prefill 阶段，返回的第二个参数为 True
    if scheduled_seqs:
        return scheduled_seqs, True

    # decode
    while self.running and num_seqs < self.max_num_seqs:
        seq = self.running.popleft() # 取出 running 队列的最左边一个请求
        # 如果不能添加这个请求，则
        while not self.block_manager.can_append(seq):
            # 如果已经有在 decode 阶段的请求，则抢占 running 队列中的最右边一个
            if self.running:
                self.preempt(self.running.pop())
            # 否则“抢占”自己，即将 seq 放回 waiting 队列的最左边
            else:
                self.preempt(seq)
                break
        # 如果能够直接将 seq 添加进 running 队列
        else:
            num_seqs += 1
            self.block_manager.may_append(seq)
            scheduled_seqs.append(seq)
    assert scheduled_seqs
    self.running.extendleft(reversed(scheduled_seqs))
    return scheduled_seqs, False

def preempt(self, seq: Sequence):
    seq.status = SequenceStatus.WAITING
    self.block_manager.deallocate(seq)
    self.waiting.appendleft(seq)

def postprocess(self, seqs: list[Sequence], token_ids: list[int]) -> list[bool]:
    for seq, token_id in zip(seqs, token_ids):
        seq.append_token(token_id) # 添加新 token
        if (not seq.ignore_eos and token_id == self.eos) or seq.num_completion_tokens == seq.max_tokens: # 生成了 eos 或达到最大 token 数量
            seq.status = SequenceStatus.FINISHED
            self.block_manager.deallocate(seq)
            self.running.remove(seq)

class Block:

    def __init__(self, block_id):
        self.block_id = block_id
        self.ref_count = 0
        self.hash = -1
        self.token_ids = []

    def update(self, hash: int, token_ids: list[int]):
        self.hash = hash
        self.token_ids = token_ids

    def reset(self):
        self.ref_count = 1
        self.hash = -1
        self.token_ids = []

def __init__(self, num_blocks: int, block_size: int):
    self.block_size = block_size
    self.blocks: list[Block] = [Block(i) for i in range(num_blocks)]
    self.hash_to_block_id: dict[int, int] = dict()
    self.free_block_ids: deque[int] = deque(range(num_blocks))
    self.used_block_ids: set[int] = set()

def can_allocate(self, seq: Sequence) -> bool:
    return len(self.free_block_ids) >= seq.num_blocks

def can_append(self, seq: Sequence) -> bool:
    return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)

def allocate(self, seq: Sequence):
    assert not seq.block_table # 没有被分配过
    h = -1                     # hash 值，默认 -1
    cache_miss = False         # 是否 cache 命中
    for i in range(seq.num_blocks):
        token_ids = seq.block(i) # 获取第 i 个 block 的 token_ids
        h = self.compute_hash(token_ids, h) if len(token_ids) == self.block_size else -1 # 如果 token_ids 可以构成一个完整的 block，则根据 token_ids 来为这个 block 计算一个 hash 值用于辨认，否则如果是不完整的块，则默认 hash 值为 -1
        block_id = self.hash_to_block_id.get(h, -1)
        if block_id == -1 or self.blocks[block_id].token_ids != token_ids:
            cache_miss = True
        if cache_miss: # 缓存不命中则新分配一个 block
            block_id = self.free_block_ids[0] # 注意这里只“引用”了第一个空闲块的 id，没有真正取出来
            block = self._allocate_block(block_id) # 在这里正式分配
        else: # 缓存命中
            seq.num_cached_tokens += self.block_size
            if block_id in self.used_block_ids: # 如果这个命中的块正在被其他请求使用
                block = self.blocks[block_id]
                block.ref_count += 1
            else: # 否则这个块是空闲的，需要通过 _allcate_block 来重新申明分配这个块
                block = self._allocate_block(block_id)
        # 如果是一个完整块，则更新 hash 值和 token_ids
        if h != -1:
            block.update(h, token_ids)
            self.hash_to_block_id[h] = block_id
        # 更新 seq 的 block table
        seq.block_table.append(block_id)

def _allocate_block(self, block_id: int) -> Block:
    block = self.blocks[block_id]
    assert block.ref_count == 0 # 分配这个块的前提是这个块当前是未被引用的
    block.reset()
    self.free_block_ids.remove(block_id)
    self.used_block_ids.add(block_id)
    return self.blocks[block_id]

@classmethod
def compute_hash(cls, token_ids: list[int], prefix: int = -1):
    h = xxhash.xxh64()
    if prefix != -1:
        h.update(prefix.to_bytes(8, "little"))
    h.update(np.array(token_ids).tobytes())
    return h.intdigest()

def deallocate(self, seq: Sequence):
    for block_id in reversed(seq.block_table): # 倒序释放
        block = self.blocks[block_id]
        block.ref_count -= 1
        if block.ref_count == 0:
            self._deallocate_block(block_id)
    seq.num_cached_tokens = 0
    seq.block_table.clear()

def _deallocate_block(self, block_id: int) -> Block:
    assert self.blocks[block_id].ref_count == 0
    self.used_block_ids.remove(block_id)
    self.free_block_ids.append(block_id)

def may_append(self, seq: Sequence):
    block_table = seq.block_table
    last_block = self.blocks[block_table[-1]]
    # 刚好多出一个 token，则新分配一个 block
    if len(seq) % self.block_size == 1:
        assert last_block.hash != -1
        block_id = self.free_block_ids[0]
        self._allocate_block(block_id)
        block_table.append(block_id)
    # 刚好满一个 block，则为该 block（seq 最后一个 block）更新 hash 值，并记录到 hash_to_block_id 表中
    elif len(seq) % self.block_size == 0:
        assert last_block.hash == -1
        token_ids = seq.block(seq.num_blocks-1)
        prefix = self.blocks[block_table[-2]].hash if len(block_table) > 1 else -1
        h = self.compute_hash(token_ids, prefix)
        last_block.update(h, token_ids)
        self.hash_to_block_id[h] = last_block.block_id
    # 否则块处于中间态，hash 值必定是 -1
    else:
        assert last_block.hash == -1

class Sampler(nn.Module):

    def __init__(self):
        super().__init__()

    @torch.compile
    def forward(self, logits: torch.Tensor, temperatures: torch.Tensor):
        logits = logits.float().div_(temperatures.unsqueeze(dim=1))
        probs = torch.softmax(logits, dim=-1)
        sample_tokens = probs.div_(torch.empty_like(probs).exponential_(1).clamp_min_(1e-10)).argmax(dim=-1)
        return sample_tokens

# __init__
self.model = Qwen3ForCausalLM(hf_config)
load_model(self.model, config.model)

# allocate_kv_cache
for module in self.model.modules():
    if hasattr(module, "k_cache") and hasattr(module, "v_cache"):
        module.k_cache = self.kv_cache[0, layer_id]
        module.v_cache = self.kv_cache[1, layer_id]
        layer_id += 1

# run_model
if is_prefill or self.enforce_eager or input_ids.size(0) > 512:
    return self.model.compute_logits(self.model(input_ids, positions))
else:
    ...
    return self.model.compute_logits(graph_vars["outputs"][:bs])

# capture_cudagraph
outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # warmup
with torch.cuda.graph(graph, self.graph_pool):
    outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # capture

# Attention.forward
k_cache, v_cache = self.k_cache, self.v_cache
if k_cache.numel() and v_cache.numel():
    store_kvcache(k, v, k_cache, v_cache, context.slot_mapping) # 存储新的 k(ey), v(alue)

# store_kvcache
def store_kvcache(key: torch.Tensor, value: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, slot_mapping: torch.Tensor):
    N, num_heads, head_dim = key.shape
    D = num_heads * head_dim
    assert key.stride(-1) == 1 and value.stride(-1) == 1
    assert key.stride(1) == head_dim and value.stride(1) == head_dim
    assert k_cache.stride(1) == D and v_cache.stride(1) == D
    assert slot_mapping.numel() == N
    store_kvcache_kernel[(N,)](key, key.stride(0), value, value.stride(0), k_cache, v_cache, slot_mapping, D)

@triton.jit
def store_kvcache_kernel(
    key_ptr,
    key_stride,
    value_ptr,
    value_stride,
    k_cache_ptr,
    v_cache_ptr,
    slot_mapping_ptr,
    D: tl.constexpr,
):
    idx = tl.program_id(0)
    slot = tl.load(slot_mapping_ptr + idx)
    if slot == -1: return
    key_offsets = idx * key_stride + tl.arange(0, D)
    value_offsets = idx * value_stride + tl.arange(0, D)
    key = tl.load(key_ptr + key_offsets)
    value = tl.load(value_ptr + value_offsets)
    cache_offsets = slot * D + tl.arange(0, D)
    tl.store(k_cache_ptr + cache_offsets, key)
    tl.store(v_cache_ptr + cache_offsets, value)

if not seq.block_table:    # warmup
    continue
for i in range(seq.num_cached_blocks, seq.num_blocks):
    start = seq.block_table[i] * self.block_size
    if i != seq.num_blocks - 1:
        end = start + self.block_size
    else:
        end = start + seq.last_block_num_tokens
    slot_mapping.extend(list(range(start, end)))