# example.py
llm = LLM(path, enforce_eager=True, tensor_parallel_size=1)
sp = SamplingParams(temperature=0.6, max_tokens=256)
outputs = llm.generate(prompts, sp)

class LLMEngine:
    def __init__(self, model, **kwargs):
        # 1. 构造 Config
        config = Config(model, **kwargs)
        # 2. 如果 TP > 1, 启动子进程 (rank 1..tp_size-1)
        for i in range(1, config.tensor_parallel_size):
            event = ctx.Event()
            p = ctx.Process(target=ModelRunner, args=(config, i, event))
            p.start()
        # 3. 主进程上初始化 rank 0 的 ModelRunner
        self.model_runner = ModelRunner(config, 0, self.events)
        # 4. tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(config.model, use_fast=True)
        config.eos = self.tokenizer.eos_token_id
        # 5. scheduler
        self.scheduler = Scheduler(config)

def add_request(self, prompt, sampling_params):
    if isinstance(prompt, str):
        prompt = self.tokenizer.encode(prompt)
    seq = Sequence(prompt, sampling_params)   # 构造 Sequence
    self.scheduler.add(seq)                   # 入 waiting 队列

def step(self):
    seqs, is_prefill = self.scheduler.schedule()                # ① 调度
    token_ids = self.model_runner.call("run", seqs, is_prefill)  # ② 跑模型
    self.scheduler.postprocess(seqs, token_ids)                 # ③ 收尾
    outputs = [(seq.seq_id, seq.completion_token_ids)
              for seq in seqs if seq.is_finished]
    return outputs, num_tokens

def generate(self, prompts, sampling_params, use_tqdm=True):
    for prompt, sp in zip(prompts, sampling_params):
        self.add_request(prompt, sp)
    while not self.is_finished():
        output, num_tokens = self.step()
        ...

class SequenceStatus(Enum):
    WAITING = auto()
    RUNNING = auto()
    FINISHED = auto()

class Sequence:
    block_size = 256
    counter = count()  # 进程内全局递增,用于恢复输出顺序

    def __init__(self, token_ids, sampling_params=SamplingParams()):
        self.seq_id = next(Sequence.counter)
        self.status = SequenceStatus.WAITING
        self.token_ids = copy(token_ids)            # prompt + 已生成
        self.last_token = token_ids[-1]             # decode 时只用这个
        self.num_tokens = len(self.token_ids)
        self.num_prompt_tokens = len(token_ids)
        self.num_cached_tokens = 0                  # prefix cache 命中数
        self.block_table = []                       # 逻辑块 → 物理块 ID
        self.temperature = sampling_params.temperature
        self.max_tokens = sampling_params.max_tokens
        self.ignore_eos = sampling_params.ignore_eos

def block(self, i):
    assert 0 <= i < self.num_blocks
    return self.token_ids[i*self.block_size: (i+1)*self.block_size]

def append_token(self, token_id):
    self.token_ids.append(token_id)
    self.last_token = token_id
    self.num_tokens += 1

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]

class Scheduler:
    def __init__(self, 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 schedule(self) -> tuple[list[Sequence], bool]:
    # 阶段 1: 优先 prefill
    scheduled_seqs = []
    while self.waiting and num_seqs < self.max_num_seqs:
        seq = self.waiting[0]
        if num_batched_tokens + len(seq) > self.max_num_batched_tokens \
           or not self.block_manager.can_allocate(seq):
            break
        self.block_manager.allocate(seq)
        num_batched_tokens += len(seq) - seq.num_cached_tokens
        seq.status = SequenceStatus.RUNNING
        self.waiting.popleft()
        self.running.append(seq)
        scheduled_seqs.append(seq)
    if scheduled_seqs:
        return scheduled_seqs, True  # is_prefill=True

    # 阶段 2: decode
    while self.running and num_seqs < self.max_num_seqs:
        seq = self.running.popleft()
        while not self.block_manager.can_append(seq):
            if self.running:
                self.preempt(self.running.pop())   # 抢占别人
            else:
                self.preempt(seq)                  # 抢占自己
                break
        else:
            num_seqs += 1
            self.block_manager.may_append(seq)
            scheduled_seqs.append(seq)
    self.running.extendleft(reversed(scheduled_seqs))  # 保持原顺序
    return scheduled_seqs, False  # is_prefill=False

def preempt(self, seq):
    seq.status = SequenceStatus.WAITING
    self.block_manager.deallocate(seq)   # 释放全部 KV block
    self.waiting.appendleft(seq)         # 放回 waiting 队头

def postprocess(self, seqs, token_ids):
    for seq, token_id in zip(seqs, token_ids):
        seq.append_token(token_id)
        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)

class Block:
    def __init__(self, block_id):
        self.block_id = block_id
        self.ref_count = 0     # 共享计数
        self.hash = -1         # -1 表示不在 prefix cache 中
        self.token_ids = []    # 冗余存储,用于 hash 命中后校验

    def update(self, hash, token_ids): ...
    def reset(self):           # 被 _allocate_block 调用
        self.ref_count = 1
        self.hash = -1
        self.token_ids = []


class BlockManager:
    def __init__(self, num_blocks, block_size):
        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()   # prefix cache 索引
        self.free_block_ids: deque[int] = deque(range(num_blocks))
        self.used_block_ids: set[int] = set()

@classmethod
def compute_hash(cls, token_ids, prefix=-1):
    h = xxhash.xxh64()
    if prefix != -1:
        h.update(prefix.to_bytes(8, "little"))   # 关键: 链入前一块的 hash
    h.update(np.array(token_ids).tobytes())      # 再 hash 当前块 token
    return h.intdigest()

def allocate(self, seq):
    assert not seq.block_table
    h = -1
    cache_miss = False
    for i in range(seq.num_blocks):
        token_ids = seq.block(i)
        h = self.compute_hash(token_ids, h) if len(token_ids) == self.block_size else -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_id = self.free_block_ids[0]
            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:                                  # 在 prefix cache 但没人用
                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 deallocate(self, seq):
    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)       # 真正放回 free 队列
    seq.num_cached_tokens = 0
    seq.block_table.clear()

def can_append(self, seq):
    # 当且仅当下一个 token 会"换块"时才需要新 block
    return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)

def may_append(self, seq):
    block_table = seq.block_table
    last_block = self.blocks[block_table[-1]]
    if len(seq) % self.block_size == 1:                       # 进入新块
        assert last_block.hash != -1                          # 旧块必须已 sealed
        block_id = self.free_block_ids[0]
        self._allocate_block(block_id)
        block_table.append(block_id)
    elif len(seq) % self.block_size == 0:                     # 刚填满一个块
        assert last_block.hash == -1                          # 必须还没 hash
        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
    else:                                                     # 普通追加
        assert last_block.hash == -1

def __init__(self, config, rank, event):
    self.config = 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

    # 1. 初始化分布式
    dist.init_process_group("nccl", "tcp://localhost:2333",
                            world_size=self.world_size, rank=rank)
    torch.cuda.set_device(rank)

    # 2. 切默认 dtype / device
    default_dtype = torch.get_default_dtype()
    torch.set_default_dtype(hf_config.torch_dtype)  # 通常是 bfloat16
    torch.set_default_device("cuda")

    # 3. 建模型 + 加载权重
    self.model = Qwen3ForCausalLM(hf_config)
    load_model(self.model, config.model)
    self.sampler = Sampler()

    # 4. warmup → 5. 分配 KV cache → 6. 捕获 CUDA Graph
    self.warmup_model()
    self.allocate_kv_cache()
    if not self.enforce_eager:
        self.capture_cudagraph()

if self.world_size > 1:
    if rank == 0:
        self.shm = SharedMemory(name="nanovllm", create=True, size=2**20)
        dist.barrier()
    else:
        dist.barrier()
        self.shm = SharedMemory(name="nanovllm")
        self.loop()           # 非 rank 0 的子进程进入 loop

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

def call(self, method_name, *args):
    if self.world_size > 1 and self.rank == 0:
        self.write_shm(method_name, *args)  # rank 0 写 SHM,通知所有子进程
    method = getattr(self, method_name, None)
    return method(*args)                   # 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)               # 跑一次 prefill
    torch.cuda.empty_cache()

def allocate_kv_cache(self):
    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   # TP 切分
    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

    # 形状: [2 (K/V), num_layers, num_blocks, block_size, num_kv_heads, head_dim]
    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)
    # 把每层 slice 挂到对应 attention 模块
    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

def prepare_prefill(self, seqs):
    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:
            continue
        for i in range(seq.num_cached_blocks, seq.num_blocks):
            start = seq.block_table[i] * self.block_size
            end = start + self.block_size if i != seq.num_blocks-1 \
                  else start + seq.last_block_num_tokens
            slot_mapping.extend(list(range(start, end)))

    # prefix cache 命中时, K 长度 > Q 长度,需要 block_table
    if cu_seqlens_k[-1] > cu_seqlens_q[-1]:
        block_tables = self.prepare_block_tables(seqs)

    # 全部转 tensor + pin_memory + non_blocking
    input_ids = torch.tensor(input_ids, dtype=torch.int64,
                             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):
    input_ids = []
    positions = []
    slot_mapping = []
    context_lens = []
    for seq in seqs:
        input_ids.append(seq.last_token)                 # 每条序列只一个 token
        positions.append(len(seq))                       # 当前位置
        context_lens.append(len(seq))                    # 上下文长度(给 FlashAttention)
        slot_mapping.append(seq.block_table[-1] * self.block_size
                            + seq.last_block_num_tokens - 1)
    ...
    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, positions, is_prefill):
    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
        for k, v in graph_vars.items():
            if k != "outputs":
                v.zero_()
        graph_vars["input_ids"][:bs] = input_ids
        graph_vars["positions"][:bs] = positions
        graph_vars["slot_mapping"][:bs] = context.slot_mapping
        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])

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

@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)
    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)
    slot = tl.load(slot_mapping_ptr + idx)
    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, value, k_cache, v_cache, slot_mapping):
    N, num_heads, head_dim = key.shape
    D = num_heads * head_dim
    store_kvcache_kernel[(N,)](key, key.stride(0), value, value.stride(0),
                               k_cache, v_cache, slot_mapping, D)

class Attention(nn.Module):
    def __init__(self, num_heads, head_dim, scale, num_kv_heads):
        super().__init__()
        ...
        self.k_cache = self.v_cache = torch.tensor([])  # 会被 ModelRunner 替换

    def forward(self, q, k, v):
        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)

        if context.is_prefill:
            if context.block_tables is not None:    # prefix cache 命中
                k, v = k_cache, v_cache             # ← 关键: K/V 来自 cache
            o = flash_attn_varlen_func(
                q, k, v,
                max_seqlen_q=context.max_seqlen_q, cu_seqlens_q=context.cu_seqlens_q,
                max_seqlen_k=context.max_seqlen_k, cu_seqlens_k=context.cu_seqlens_k,
                softmax_scale=self.scale, causal=True,
                block_table=context.block_tables,
            )
        else:    # decode
            o = flash_attn_with_kvcache(
                q.unsqueeze(1), k_cache, v_cache,
                cache_seqlens=context.context_lens,
                block_table=context.block_tables,
                softmax_scale=self.scale, causal=True,
            )
        o = o.view(-1, self.num_heads * self.head_dim)
        return o

def forward(self, positions, hidden_states, residual):
    if residual is None:
        # 第一层: 没有 residual
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
    else:
        # 后续层: 把残差加和 norm 融合
        hidden_states, residual = self.input_layernorm(hidden_states, residual)
    hidden_states = self.self_attn(positions, hidden_states)
    hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
    hidden_states = self.mlp(hidden_states)
    return hidden_states, residual

@torch.compile
def add_rms_forward(self, x, residual):
    orig_dtype = x.dtype
    x = x.to(torch.float32).add_(residual.to(torch.float32))  # 融合 1
    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

def forward(self, positions, hidden_states):
    qkv = self.qkv_proj(hidden_states)
    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
    q_by_head = q.view(-1, self.num_heads, self.head_dim)
    q_by_head = self.q_norm(q_by_head)               # Qwen3 特有
    q = q_by_head.view(q.shape)
    k_by_head = k.view(-1, self.num_kv_heads, self.head_dim)
    k_by_head = self.k_norm(k_by_head)               # Qwen3 特有
    k = k_by_head.view(k.shape)
    q, k = self.rotary_emb(positions, q, k)
    o = self.attn(q, k, v)
    output = self.o_proj(o)
    return output

class Qwen3MLP(nn.Module):
    def forward(self, x):
        gate_up = self.gate_up_proj(x)   # 一次性输出 [gate, up]
        x = self.act_fn(gate_up)         # silu(gate) * up
        x = self.down_proj(x)
        return x

@torch.compile
def forward(self, x):
    x, y = x.chunk(2, -1)
    return F.silu(x) * y

class Qwen3ForCausalLM(nn.Module):
    packed_modules_mapping = {
        "q_proj": ("qkv_proj", "q"),
        "k_proj": ("qkv_proj", "k"),
        "v_proj": ("qkv_proj", "v"),
        "gate_proj": ("gate_up_proj", 0),
        "up_proj": ("gate_up_proj", 1),
    }

    def forward(self, input_ids, positions):
        return self.model(input_ids, positions)

    def compute_logits(self, hidden_states):
        return self.lm_head(hidden_states)

class ColumnParallelLinear(LinearBase):
    def __init__(self, input_size, output_size, bias=False):
        super().__init__(input_size, output_size, 0)   # tp_dim=0 表示按行切
        self.output_size_per_partition = divide(output_size, self.tp_size)
        self.weight = nn.Parameter(torch.empty(self.output_size_per_partition, self.input_size))
        ...

    def weight_loader(self, param, loaded_weight):
        param_data = param.data
        shard_size = param_data.size(self.tp_dim)
        start_idx = self.tp_rank * shard_size
        loaded_weight = loaded_weight.narrow(self.tp_dim, start_idx, shard_size)
        param_data.copy_(loaded_weight)

    def forward(self, x):
        return F.linear(x, self.weight, self.bias)   # 不需要 all-reduce

class RowParallelLinear(LinearBase):
    def __init__(self, input_size, output_size, bias=False):
        super().__init__(input_size, output_size, 1)   # tp_dim=1 表示按列切
        self.input_size_per_partition = divide(input_size, self.tp_size)
        self.weight = nn.Parameter(torch.empty(self.output_size, self.input_size_per_partition))
        ...

    def forward(self, x):
        y = F.linear(x, self.weight, self.bias if self.tp_rank == 0 else None)
        if self.tp_size > 1:
            dist.all_reduce(y)
        return y

class QKVParallelLinear(ColumnParallelLinear):
    def __init__(self, hidden_size, head_size, total_num_heads, total_num_kv_heads=None, bias=False):
        ...
        output_size = (total_num_heads + 2 * total_num_kv_heads) * head_size
        super().__init__(hidden_size, output_size, bias=bias)

    def weight_loader(self, param, loaded_weight, loaded_shard_id):
        # loaded_shard_id ∈ {"q", "k", "v"}
        if loaded_shard_id == "q":
            shard_offset = 0
            shard_size = self.num_heads * self.head_size
        elif loaded_shard_id == "k":
            shard_offset = self.num_heads * self.head_size
            shard_size = self.num_kv_heads * self.head_size
        else:  # "v"
            shard_offset = self.num_heads * self.head_size + self.num_kv_heads * self.head_size
            shard_size = self.num_kv_heads * self.head_size
        param_data = param.data.narrow(self.tp_dim, shard_offset, shard_size)
        loaded_weight = loaded_weight.chunk(self.tp_size, self.tp_dim)[self.tp_rank]
        param_data.copy_(loaded_weight)

class VocabParallelEmbedding(nn.Module):
    def forward(self, x):
        if self.tp_size > 1:
            mask = (x >= self.vocab_start_idx) & (x < self.vocab_end_idx)
            x = mask * (x - self.vocab_start_idx)
        y = F.embedding(x, self.weight)
        if self.tp_size > 1:
            y = mask.unsqueeze(1) * y
            dist.all_reduce(y)
        return y

class ParallelLMHead(VocabParallelEmbedding):
    def forward(self, x):
        context = get_context()
        if context.is_prefill:
            # 只取每条序列最后一个 token 的 hidden state
            last_indices = context.cu_seqlens_q[1:] - 1
            x = x[last_indices].contiguous()
        logits = F.linear(x, self.weight, self.bias)
        if self.tp_size > 1:
            all_logits = [torch.empty_like(logits) for _ in range(self.tp_size)] \
                         if self.tp_rank == 0 else None
            dist.gather(logits, all_logits, 0)
            logits = torch.cat(all_logits, -1) if self.tp_rank == 0 else None
        return logits

class Sampler(nn.Module):
    def forward(self, logits, temperatures):
        logits = logits.to(torch.float)
        greedy_tokens = logits.argmax(dim=-1)
        logits.div_(temperatures.unsqueeze(dim=1))
        probs = torch.softmax(logits, dim=-1, dtype=torch.float)
        sample_tokens = probs.div_(torch.empty_like(probs).exponential_(1) + epsilon).argmax(dim=-1)
        return torch.where(temperatures == 0, greedy_tokens, sample_tokens)

def default_weight_loader(param, loaded_weight):
    param.data.copy_(loaded_weight)

def load_model(model, path):
    packed_modules_mapping = getattr(model, "packed_modules_mapping", {})
    for file in glob(os.path.join(path, "*.safetensors")):
        with safe_open(file, "pt", "cpu") as f:
            for weight_name in f.keys():
                for k in packed_modules_mapping:
                    if k in weight_name:
                        v, shard_id = packed_modules_mapping[k]
                        param_name = weight_name.replace(k, v)
                        param = model.get_parameter(param_name)
                        weight_loader = getattr(param, "weight_loader")
                        weight_loader(param, f.get_tensor(weight_name), shard_id)
                        break
                else:
                    param = model.get_parameter(weight_name)
                    weight_loader = getattr(param, "weight_loader", default_weight_loader)
                    weight_loader(param, f.get_tensor(weight_name))

@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

_CONTEXT = Context()

def get_context():
    return _CONTEXT

def set_context(is_prefill, ...):
    global _CONTEXT
    _CONTEXT = Context(is_prefill, ...)

def reset_context():
    global _CONTEXT
    _CONTEXT = Context()

class RotaryEmbedding(nn.Module):
    def __init__(self, head_size, rotary_dim, max_position_embeddings, base):
        ...
        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()
        self.register_buffer("cos_sin_cache", torch.cat((cos, sin), dim=-1), persistent=False)

    @torch.compile
    def forward(self, positions, query, key):
        cos_sin = self.cos_sin_cache[positions]
        cos, sin = cos_sin.chunk(2, dim=-1)
        query = apply_rotary_emb(query.view(num_tokens, -1, head_size), cos, sin).view(query_shape)
        key = apply_rotary_emb(key.view(num_tokens, -1, head_size), cos, sin).view(key_shape)
        return query, key

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

@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     # 必须能被 256 整除
        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