def apply_rotary_emb(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
) -> torch.Tensor:
    x1, x2 = torch.chunk(x.float(), 2, dim=-1)
    y1 = x1 * cos - x2 * sin
    y2 = x2 * cos + x1 * sin
    return torch.cat((y1, y2), dim=-1).to(x.dtype)

inv_freq = 1.0 / (base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim))
t = torch.arange(max_position_embeddings, dtype=torch.float)
freqs = torch.einsum("i,j->ij", t, inv_freq)
cos = freqs.cos()
sin = freqs.sin()
cache = torch.cat((cos, sin), dim=-1).unsqueeze_(1)
self.register_buffer("cos_sin_cache", cache, persistent=False)

@torch.compile
def rms_forward(self, x: torch.Tensor) -> torch.Tensor:
    orig_dtype = x.dtype
    x = x.float()
    var = x.pow(2).mean(dim=-1, keepdim=True)
    x.mul_(torch.rsqrt(var + self.eps))
    x = x.to(orig_dtype).mul_(self.weight)
    return x

# 残差融合的 RMSNorm
# 输入：x, residual
# 输出: RMSNorm(x + residual), x + residual
# 返回的 residual 将作为下一层（如 FFN）输入，并再次参与残差连接。
@torch.compile
def add_rms_forward(self, x: torch.Tensor, residual: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    orig_dtype = x.dtype
    x = x.float().add_(residual.float())
    residual = x.to(orig_dtype)
    var = x.pow(2).mean(dim=-1, keepdim=True)
    x.mul_(torch.rsqrt(var + self.eps))
    x = x.to(orig_dtype).mul_(self.weight)
    return x, residual

# 输入形状需要从（batch_size, seq_len, num_heads, head_dim）转换为（total_q, num_heads, head_dim），total_q是所有序列长度之和
# 输出的o是(total_q, num_heads, head_dim)
o = flash_attn_varlen_func(q, # (total_q, num_heads, head_dim)
k, # (total_k, num_kv_heads, head_dim)
v, # (total_v, num_kv_heads, head_dim)
max_seqlen_q=context.max_seqlen_q,  # max sequence length of query
cu_seqlens_q=context.cu_seqlens_q, # (batch_size + 1,) cumulative sequence lengths of query
max_seqlen_k=context.max_seqlen_k, # max sequence length of key and value
cu_seqlens_k=context.cu_seqlens_k, # (batch_size + 1,) cumulative sequence lengths of key and value
softmax_scale=self.scale,  # head_dim ** -0.5
causal=True, # whether to apply causal mask
)

# 输出的o是(total_q, num_heads, head_dim)
o = flash_attn_varlen_func(q, # (total_q, num_heads, head_dim)
k=k_cache, # (num_blocks, block_size, num_kv_heads, head_dim)
v=v_cache, # (num_blocks, block_size, num_kv_heads, head_dim)
max_seqlen_q=context.max_seqlen_q,  # max sequence length of query
cu_seqlens_q=context.cu_seqlens_q, # (batch_size + 1,) cumulative sequence lengths of query
max_seqlen_k=context.max_seqlen_k, # max sequence length of key and value
cu_seqlens_k=context.cu_seqlens_k, # (batch_size + 1,) cumulative sequence lengths of key and value
softmax_scale=self.scale,  # head_dim ** -0.5
causal=True, # whether to apply causal mask
block_table=context.block_tables, # (batch_size, max_num_blocks_per_seq). max_num_blocks_per_seq表示当前 Batch 中最长的那个序列所占用的 Block 数量。对于短序列，后面的空位会用 -1 进行 Padding。
)

# o是(batch_size, 1, num_heads, head_dim)
o = flash_attn_with_kvcache(q.unsqueeze(1), # (batch_size, 1, num_heads, head_dim), 当前 Batch 中每个序列生成的最新 Token。注意第 1 维是 Sequence Length，Decode 阶段始终为 1。 
k_cache, # (num_blocks, block_size, num_kv_heads, head_dim)
v_cache, # (num_blocks, block_size, num_kv_heads, head_dim)
cache_seqlens=context.context_lens, # (batch_size + 1,) cumulative sequence lengths of key and value
block_table=context.block_tables, # (batch_size, max_num_blocks_per_seq). max_num_blocks_per_seq表示当前 Batch 中最长的那个序列所占用的 Block 数量。对于短序列，后面的空位会用 -1 进行 Padding。
softmax_scale=self.scale,  # head_dim ** -0.5
causal=True, # whether to apply causal mask
)

class Sequence:
    block_size = 256
    counter = count()

    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
        # 这个时候的block table是纯在CPU上的，在context上会组合batch内sequence的block table变成大小b（atch_size, max_num_blocks_per_seq). max_num_blocks_per_seq表示当前 Batch 中最长的那个序列所占用的 Block 数量。对于短序列，后面的空位会用 -1 进行 Padding。
        self.block_table = []
        self.temperature = sampling_params.temperature
        self.max_tokens = sampling_params.max_tokens
        self.ignore_eos = sampling_params.ignore_eos

    @property
    def num_cached_blocks(self):
        return self.num_cached_tokens // self.block_size
    
    @property
    def num_blocks(self):
	    # 当前这个序列需要多少个block来存储
        return (self.num_tokens + self.block_size - 1) // self.block_size
    
    @property
    def last_block_num_tokens(self):
        return self.num_tokens - (self.num_blocks - 1) * self.block_size
    
    def block(self, i):
        assert 0 <= i < self.num_blocks
        # Python 的 list 切片在右边界（甚至左边界）超出范围时，不会报错，而是自动截断到合法范围。
        return self.token_ids[i * self.block_size: (i+1) * self.block_size]

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 step(self):
    # 每个step，从scheduler里面得到下一轮迭代要跑的sequence, 可能是prefill, 也可能是decode
    seqs, is_prefill = self.scheduler.schedule()
    token_ids = self.model_runner.call("run", seqs, is_prefill)
    # postprocess主要是把token_ids添加到sequence中，并判断是否达到eos或max_tokens
    self.scheduler.postprocess(seqs, token_ids)
    # 返回所有已经完成的sequence的token_ids
    outputs = [(seq.seq_id, seq.completion_token_ids) for seq in seqs if seq.is_finished]
    # 返回这轮生成的token数量
    num_tokens = sum(len(seq) for seq in seqs) if is_prefill else -len(seqs)
    return outputs, num_tokens

def generate(self, prompts: list[str] | list[list[int]], sampling_params: SamplingParams | list[SamplingParams],
use_tqdm: bool = True) -> list[str]:
    ...
    for prompt, sp in zip(prompts, sampling_params):
        self.add_request(prompt, sp)
    outputs = {}
    prefill_throughput = decode_throughput = 0.
    while not self.is_finished():
        t = perf_counter()
        output, num_tokens = self.step()
        ...
    return outputs

    def schedule(self) -> tuple[list[Sequence], bool]:
        # 这个函数的目标是选出一批 Sequence 来执行一次 forward。它优先服务 Prefill（新来的或被抢占的），如果没有 Prefill 任务，再服务 Decode（正在生成的）。
        # prefill
        scheduled_seqs = []
        num_seqs = 0
        # 当前的这个batch中已经有的token数量
        num_batched_tokens = 0

        # waiting的意思是还没有进行过prefill或者被抢占的
        # self.waiting队列里面既有刚进来的全新请求，也有之前运行一半因为显存不够被抢占的
        while self.waiting and num_seqs < self.max_num_seqs:
            # 从waiting的队列中取出第一个sequence
            seq = self.waiting[0]
            # 如果当前的这个batch中已经有的token数量加上这个sequence的长度超过了max_num_batched_tokens，或者当前free_block_ids数量不足以分配给这个sequence，则break
            if num_batched_tokens + len(seq) > self.max_num_batched_tokens or not self.block_manager.can_allocate(seq):
                break
            num_seqs += 1
            # 给这个sequence分配更多的block用于用于存储KV
            # 注意如果之前这个sequence被抢占过，那么有可能部分前缀的KV没有被覆盖，可以重新复用, 这部分就不需要做prefill了
            self.block_manager.allocate(seq)
            # 对没有被缓存的token做prefill，说明支持增量prefill
            num_batched_tokens += len(seq) - seq.num_cached_tokens
            seq.status = SequenceStatus.RUNNING
            self.waiting.popleft()
            self.running.append(seq)
            scheduled_seqs.append(seq)
        # 如果成功调度了至少一个 Prefill 任务，直接返回 (seqs, True)。这意味着一旦有 Prefill 任务，这一轮就只跑 Prefill，不跑 Decode（因为两者计算特征不同，很少混着跑）。
        if scheduled_seqs:
            return scheduled_seqs, True

        # decode
        while self.running and num_seqs < self.max_num_seqs:
            seq = self.running.popleft()
            # 主要检查能不能再生成一个 Token. 因为生成一个新 Token 需要占用新的 KV Cache block或者占用当前 Block 的剩余空间
            while not self.block_manager.can_append(seq):
                # “受害者选择”调度策略
                # 如果显存不足，尝试腾出空间
                if self.running:
                    # 如果队列里还有其他正在运行的序列，就执行 “牺牲他人” 策略。它会把队列尾部的序列弹出, 并对其执行挂起self.preempt(...)
                    self.preempt(self.running.pop())
                else:
                    # 如果队列空了（只剩当前这个 seq 了）但显存还是不够，那就说明连运行这一个序列的资源都不够。此时只能 “牺牲自己” 
                    self.preempt(seq)
                    break
            else:
                num_seqs += 1
                # 添加一个token，如果当前最后一个block还有剩余空间，则直接添加，否则需要分配新的block
                self.block_manager.may_append(seq)
                scheduled_seqs.append(seq)
        assert scheduled_seqs
        self.running.extendleft(reversed(scheduled_seqs))
        return scheduled_seqs, False

def _deallocate_block(self, block_id: int) -> Block:
    # 关键点：虽然ref_count 降为 0，Block ID 会被放入 free_block_ids（空闲队列），但是 hash_to_block_id里依然保留着记录（除非被别的请求覆盖了）。
    assert self.blocks[block_id].ref_count == 0
    self.used_block_ids.remove(block_id)
    self.free_block_ids.append(block_id)

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)
    # 在当前nano-vllm的实现中，如果被抢占了，那么num_cached_tokens会被设置为0
    seq.num_cached_tokens = 0
    seq.block_table.clear() 

def allocate(self, seq: Sequence):
    # Prefill 阶段分配
    # 当一个新的 Sequence（或者被抢占回来的）准备开始 Prefill 时调用。它的目标是为整个 Sequence 分配足够的 Block，并尽可能复用已有的 Block（Prefix Caching）
    assert not seq.block_table
    h = -1
    cache_miss = False
    for i in range(seq.num_blocks):
        token_ids = seq.block(i)
        # 如果当前的 Block 是填满的，就基于它的内容和前一个 Block 的哈希值计算新的哈希；否则（如果是未填满的尾部 Block），不计算哈希（设为 -1）, 也不做缓存
        h = self.compute_hash(token_ids, h) if len(token_ids) == self.block_size else -1
        # 系统算出 Hash，去查表。
        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_id = self.free_block_ids[0]
            block = self._allocate_block(block_id)
        else:
            # 如果发现发现 hash_to_block_id 里有记录(block_id!=-1), 并且刚好就是当前的token_ids，那么就可以直接复用这个block，就可以直接复用了
            seq.num_cached_tokens += self.block_size
            # 如果block_id已经在used_block_ids中，则说明这个block被其他sequence正在使用，需要增加ref_count
            if block_id in self.used_block_ids:
                block = self.blocks[block_id]
                block.ref_count += 1
            # 如果block_id不在used_block_ids中，则说明这个block是空闲的，但是之前的记录没有被覆盖，只需要形式上重新分配一次即可，但是num_cached_tokens还是可以增加
            else:
                block = self._allocate_block(block_id)
        if h != -1:
            block.update(h, token_ids)
            self.hash_to_block_id[h] = block_id
        seq.block_table.append(block_id)

    def allocate_kv_cache(self):
        config = self.config
        hf_config = config.hf_config
        free, total = torch.cuda.mem_get_info()
        used = total - free # GPU 上当前已使用的总内存（包括所有进程和 PyTorch 的内存池）
        peak = torch.cuda.memory_stats()['allocated_bytes.all.peak']
        current = torch.cuda.memory_stats()['allocated_bytes.all.current'] # PyTorch 当前实际分配的内存（不包括内存池中的空闲块）
        num_kv_heads = hf_config.num_key_value_heads // self.world_size
        head_dim = getattr(hf_config, "head_dim", hf_config.hidden_size // hf_config.num_attention_heads)
        block_bytes = 2 * hf_config.num_hidden_layers * self.block_size * num_kv_heads * 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.empty(2, hf_config.num_hidden_layers, config.num_kvcache_blocks, self.block_size, num_kv_heads, 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

@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)
    
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)
    

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
        context = get_context()
        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)

    def prepare_prefill(self, seqs: list[Sequence]):
		...
        for seq in seqs:
            seqlen = len(seq)
            ...
            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)))

    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 # 只包含新 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)))
        # 这个batch里面存在一个sequence有cache，就是cu_seqlens_k[-1] > cu_seqlens_q[-1]，需要使用prefix cache
        if cu_seqlens_k[-1] > cu_seqlens_q[-1]:
            block_tables = self.prepare_block_tables(seqs)
        # https://pytorch-tutorial.web.app/intermediate/pinmem_nonblock.html
        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) 
        # (batch_size), 每个序列最新生成的 token
        positions = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True) 
        # (batch_size),每个序列最新token的位置索引（用于 RoPE）
        slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) 
        # (batch_size)，每个序列的完整长度（attention 需要关注的范围）
        context_lens = torch.tensor(context_lens, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) 
        # (batch_size)，新 token 的 KV cache 存储位置
        block_tables = self.prepare_block_tables(seqs) 
        # 因为 decode 阶段的每个新 token 的 Query 需要关注所有之前的 KV cache，
        # 必须通过 block_tables 告诉 attention kernel 
        # 如何在 paged memory 中找到这些缓存的 K、V。
        set_context(False, slot_mapping=slot_mapping, context_lens=context_lens, block_tables=block_tables)
        return input_ids, positions

    @torch.inference_mode()
    def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill:bool):
        return self.model.compute_logits(self.model(input_ids, positions))
    
    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

import torch
from torch import nn

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))
        # softmax归一化
        probs = torch.softmax(logits, dim=-1)
        # Gumbel-Max采样
        # G_i ~ Exponential(1)  (指数分布)
        # 任何分类分布除以指数分布采样，等价于直接采样
        # sample = argmax_i (p_i / G_i) = argmax_i (log(p_i) - G_i)
        # 虽然采样本身不可微，但可以通过 Gumbel-Softmax 的连续松弛实现可微, 同时数学上等价于从分类分布中采样
        sample_tokens = probs.div_(torch.empty_like(probs).exponential_(1).clamp_min_(1e-10)).argmax(dim=-1)
        return sample_tokens