1
from nanovllm.llm import LLM
2
from nanovllm.sampling_params import SamplingParams

1
@dataclass(slots=True)
2
class Config:
3
    model: str                              # 模型文件夹路径
4
    max_num_batched_tokens: int = 16384     # 一批最多处理多少个 token
5
    max_num_seqs: int = 512                 # 一批最多同时处理多少个请求
6
    max_model_len: int = 4096               # 模型最大上下文长度
7
    gpu_memory_utilization: float = 0.9     # 最多使用 GPU 显存的 90%
8
    tensor_parallel_size: int = 1           # 用几张 GPU（1=单卡）
9
    enforce_eager: bool = False             # 是否禁用 CUDA Graph
10
    hf_config: AutoConfig | None = None     # HuggingFace 的模型配置
11
    eos: int = -1                           # 结束符的 token ID
12
    kvcache_block_size: int = 256           # 每个 KV-Cache 块的大小
13
    num_kvcache_blocks: int = -1            # KV-Cache 块的数量（自动计算）

1
def __post_init__(self):
2
    assert os.path.isdir(self.model)                         # 1. 确保模型目录存在
3
    assert self.kvcache_block_size % 256 == 0                # 2. 块大小必须是 256 的倍数
4
    assert 1 <= self.tensor_parallel_size <= 8               # 3. GPU 数量在 1-8 之间
5
    self.hf_config = AutoConfig.from_pretrained(self.model)  # 4. 读取模型配置文件
6
    self.max_model_len = min(self.max_model_len, self.hf_config.max_position_embeddings)  # 5. 限制最大长度

1
@dataclass(slots=True)
2
class SamplingParams:
3
    temperature: float = 1.0    # 温度：控制随机性（越低越确定）
4
    max_tokens: int = 64        # 最多生成多少个 token
5
    ignore_eos: bool = False    # 是否忽略结束符（一直生成到 max_tokens）
6

7
    def __post_init__(self):
8
        assert self.temperature > 1e-10, "greedy sampling is not permitted"

1
from nanovllm.engine.llm_engine import LLMEngine
2

3
class LLM(LLMEngine):
4
    pass

1
class SequenceStatus(Enum):
2
    WAITING = auto()        # 等待中（还没开始算）
3
    RUNNING = auto()        # 运行中（正在处理）
4
    FINISHED = auto()       # 已完成

1
class Sequence:
2
    block_size = 256       # 类变量：每个块的大小
3
    counter = count()      # 类变量：自增计数器（从0开始，每次创建序列+1）
4

5
    def __init__(self, token_ids, sampling_params):
6
        self.seq_id = next(Sequence.counter)     # 自动分配唯一 ID
7
        self.status = SequenceStatus.WAITING     # 初始状态：等待
8
        self.token_ids = copy(token_ids)         # 深拷贝 token 列表
9
        self.last_token = token_ids[-1]          # 最后一个 token（decode 时需要）
10
        self.num_tokens = len(self.token_ids)    # 当前总 token 数
11
        self.num_prompt_tokens = len(token_ids)  # 输入 token 数（不含生成）
12
        self.num_cached_tokens = 0               # 被前缀缓存命中的 token 数
13
        self.num_scheduled_tokens = 0            # 本轮计划处理的 token 数
14
        self.is_prefill = True                   # 是否处于预填充阶段
15
        self.block_table = []                    # KV-Cache 块表
16
        self.temperature = ...                   # 从采样参数复制
17
        self.max_tokens = ...
18
        self.ignore_eos = ...

1
@property
2
def num_blocks(self):
3
    return (self.num_tokens + self.block_size - 1) // self.block_size

1
def block(self, i):
2
    return self.token_ids[i*self.block_size: (i+1)*self.block_size]

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

1
def __getstate__(self):
2
    last_state = self.last_token if not self.is_prefill else self.token_ids
3
    return (self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens,
4
            self.num_scheduled_tokens, self.block_table, last_state)
5

6
def __setstate__(self, state):
7
    # 从序列化状态中恢复
8
    ...

1
class Block:
2
    def __init__(self, block_id):
3
        self.block_id = block_id    # 块的编号
4
        self.ref_count = 0          # 引用计数：几个序列在用这个块
5
        self.hash = -1              # 内容的哈希值（用于前缀缓存匹配）
6
        self.token_ids = []         # 块中存储的 token ID

1
class BlockManager:
2
    def __init__(self, num_blocks, block_size):
3
        self.blocks = [Block(i) for i in range(num_blocks)]     # 所有块
4
        self.hash_to_block_id = dict()                          # 哈希 → 块ID 映射
5
        self.free_block_ids = deque(range(num_blocks))          # 空闲块队列
6
        self.used_block_ids = set()                             # 已使用块集合

1
@classmethod
2
def compute_hash(cls, token_ids, prefix=-1):
3
    h = xxhash.xxh64()                          # 创建 xxHash 对象（很快的哈希算法）
4
    if prefix != -1:
5
        h.update(prefix.to_bytes(8, "little"))  # 把前一个块的哈希也混入
6
    h.update(np.array(token_ids).tobytes())     # 把 token ID 数组转成字节混入
7
    return h.intdigest()                        # 返回 64 位哈希值

1
def can_allocate(self, seq):
2
    """检查能否为一个序列分配 KV-Cache，返回前缀命中的块数"""
3
    h = -1
4
    num_cached_blocks = 0
5
    num_new_blocks = seq.num_blocks
6
    for i in range(seq.num_blocks - 1):
7
        token_ids = seq.block(i)
8
        h = self.compute_hash(token_ids, h)
9
        block_id = self.hash_to_block_id.get(h, -1)
10
        if block_id == -1 or self.blocks[block_id].token_ids != token_ids:
11
            break                              # 前缀匹配中断
12
        num_cached_blocks += 1
13
        if block_id in self.used_block_ids:
14
            num_new_blocks -= 1                # 这个块已经被其他序列用了，不占用新空间
15
    if len(self.free_block_ids) < num_new_blocks:
16
        return -1                              # 空闲块不够，无法分配
17
    return num_cached_blocks

1
def allocate(self, seq, num_cached_blocks):
2
    """为序列分配 KV-Cache 块"""
3
    h = -1
4
    # 第一步：处理前缀命中的块（增加引用计数）
5
    for i in range(num_cached_blocks):
6
        token_ids = seq.block(i)
7
        h = self.compute_hash(token_ids, h)
8
        block_id = self.hash_to_block_id[h]
9
        block = self.blocks[block_id]
10
        if block_id in self.used_block_ids:
11
            block.ref_count += 1               # 共享块，引用+1
12
        else:
13
            block.ref_count = 1                # 之前空闲的块，标记为使用
14
            self.free_block_ids.remove(block_id)
15
            self.used_block_ids.add(block_id)
16
        seq.block_table.append(block_id)
17
    # 第二步：为新块分配空间
18
    for i in range(num_cached_blocks, seq.num_blocks):
19
        seq.block_table.append(self._allocate_block())
20
    seq.num_cached_tokens = num_cached_blocks * self.block_size

1
def deallocate(self, seq):
2
    """释放序列占用的 KV-Cache 块"""
3
    for block_id in reversed(seq.block_table):
4
        block = self.blocks[block_id]
5
        block.ref_count -= 1
6
        if block.ref_count == 0:
7
            self._deallocate_block(block_id)   # 没有其他人引用，真正释放
8
    seq.num_cached_tokens = 0
9
    seq.block_table.clear()

1
def can_append(self, seq):
2
    """检查序列是否需要并能够分配一个新块"""
3
    return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)

1
def may_append(self, seq):
2
    """如果需要，为序列追加一个新块"""
3
    if len(seq) % self.block_size == 1:
4
        seq.block_table.append(self._allocate_block())

1
def hash_blocks(self, seq):
2
    """计算并存储序列块的哈希值（为前缀缓存做准备）"""
3
    start = seq.num_cached_tokens // self.block_size
4
    end = (seq.num_cached_tokens + seq.num_scheduled_tokens) // self.block_size
5
    if start == end: return
6
    h = self.blocks[seq.block_table[start - 1]].hash if start > 0 else -1
7
    for i in range(start, end):
8
        block = self.blocks[seq.block_table[i]]
9
        token_ids = seq.block(i)
10
        h = self.compute_hash(token_ids, h)
11
        block.update(h, token_ids)                     # 更新块的哈希和内容
12
        self.hash_to_block_id[h] = block.block_id      # 注册到哈希表

1
class Scheduler:
2
    def __init__(self, config):
3
        self.max_num_seqs = config.max_num_seqs          # 最多同时处理多少请求
4
        self.max_num_batched_tokens = config.max_num_batched_tokens  # 一批最多多少 token
5
        self.eos = config.eos                            # 结束符 token ID
6
        self.block_size = config.kvcache_block_size
7
        self.block_manager = BlockManager(...)
8
        self.waiting = deque()    # 等待队列（还没开始或被打断的请求）
9
        self.running = deque()    # 运行队列（正在处理的请求）

1
def schedule(self) -> tuple[list[Sequence], bool]:

1
while self.waiting and len(scheduled_seqs) < self.max_num_seqs:
2
    seq = self.waiting[0]
3
    remaining = self.max_num_batched_tokens - num_batched_tokens
4
    if remaining == 0:
5
        break                                            # 1. 本轮 token 预算用完
6
    if not seq.block_table:
7
        num_cached_blocks = self.block_manager.can_allocate(seq)
8
        if num_cached_blocks == -1:
9
            break                                        # 2. KV-Cache 不够
10
        num_tokens = seq.num_tokens - num_cached_blocks * self.block_size
11
    else:
12
        num_tokens = seq.num_tokens - seq.num_cached_tokens  # 被打断后继续
13
    if remaining < num_tokens and scheduled_seqs:        # 3. 只对第一个序列做 chunked prefill
14
        break
15
    if not seq.block_table:
16
        self.block_manager.allocate(seq, num_cached_blocks)
17
    seq.num_scheduled_tokens = min(num_tokens, remaining)
18
    num_batched_tokens += seq.num_scheduled_tokens
19
    if seq.num_cached_tokens + seq.num_scheduled_tokens == seq.num_tokens:
20
        seq.status = SequenceStatus.RUNNING              # 该序列 prefill 完成
21
        self.waiting.popleft()
22
        self.running.append(seq)
23
    scheduled_seqs.append(seq)

1
while self.running and len(scheduled_seqs) < self.max_num_seqs:
2
    seq = self.running.popleft()
3
    while not self.block_manager.can_append(seq):
4
        if self.running:
5
            self.preempt(self.running.pop())    # 内存不够，抢占最后一个 running 序列
6
        else:
7
            self.preempt(seq)                   # 只有它自己，只能抢占自己
8
            break
9
    else:
10
        seq.num_scheduled_tokens = 1
11
        seq.is_prefill = False
12
        self.block_manager.may_append(seq)
13
        scheduled_seqs.append(seq)

1
def preempt(self, seq):
2
    seq.status = SequenceStatus.WAITING
3
    seq.is_prefill = True
4
    self.block_manager.deallocate(seq)    # 释放 KV-Cache
5
    self.waiting.appendleft(seq)          # 放回等待队列头部（优先处理）

1
def postprocess(self, seqs, token_ids, is_prefill):
2
    for seq, token_id in zip(seqs, token_ids):
3
        self.block_manager.hash_blocks(seq)             # 给刚算完的块打哈希标签
4
        seq.num_cached_tokens += seq.num_scheduled_tokens
5
        seq.num_scheduled_tokens = 0
6
        if is_prefill and seq.num_cached_tokens < seq.num_tokens:
7
            continue                                    # chunked prefill 未完成，不追加 token
8
        seq.append_token(token_id)                      # 追加生成的新 token
9
        if (not seq.ignore_eos and token_id == self.eos) or \
10
           seq.num_completion_tokens == seq.max_tokens:
11
            seq.status = SequenceStatus.FINISHED         # 达到结束条件
12
            self.block_manager.deallocate(seq)
13
            self.running.remove(seq)

1
class ModelRunner:
2
    def __init__(self, config, rank, event):
3
        # 1. 初始化分布式通信
4
        dist.init_process_group("nccl", "tcp://localhost:2333", ...)
5
        torch.cuda.set_device(rank)
6

7
        # 2. 创建模型
8
        self.model = Qwen3ForCausalLM(hf_config)
9
        load_model(self.model, config.model)    # 加载权重
10

11
        # 3. 准备
12
        self.sampler = Sampler()
13
        self.warmup_model()           # 预热：跑一遍空数据
14
        self.allocate_kv_cache()      # 分配 KV-Cache 显存
15
        if not self.enforce_eager:
16
            self.capture_cudagraph()  # 录制 CUDA Graph

1
def allocate_kv_cache(self):
2
    # 计算还剩多少可用显存
3
    free, total = torch.cuda.mem_get_info()
4
    used = total - free
5
    peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"]
6
    current = torch.cuda.memory_stats()["allocated_bytes.all.current"]
7

8
    # 计算每个块占多少字节
9
    num_kv_heads = hf_config.num_key_value_heads // self.world_size
10
    head_dim = hf_config.hidden_size // hf_config.num_attention_heads
11
    block_bytes = 2 * num_layers * block_size * num_kv_heads * head_dim * dtype_size
12
    # 2 = K 和 V 两份，每个 token 每层都有各自的 K、V
13

14
    # 计算可以分配多少块
15
    config.num_kvcache_blocks = int(
16
        total * gpu_memory_utilization - used - peak + current
17
    ) // block_bytes
18

19
    # 创建 KV-Cache 张量（形状：(2, 层数, 块数, 块大小, KV头数, 头维度)）
20
    self.kv_cache = torch.empty(2, num_layers, num_blocks, block_size, num_kv_heads, head_dim)
21

22
    # 把缓存分配给每个 Attention 层
23
    for module in self.model.modules():
24
        if hasattr(module, "k_cache") and hasattr(module, "v_cache"):
25
            module.k_cache = self.kv_cache[0, layer_id]
26
            module.v_cache = self.kv_cache[1, layer_id]

1
def prepare_prefill(self, seqs):
2
    input_ids = []          # 所有待计算的 token ID
3
    positions = []          # 每个 token 在序列中的位置
4
    cu_seqlens_q = [0]      # 累积查询长度（用于 varlen attention）
5
    cu_seqlens_k = [0]      # 累积键长度
6
    slot_mapping = []       # 每个 token 的 KV-Cache 位置
7

8
    for seq in seqs:
9
        start = seq.num_cached_tokens           # 之前缓存了多少
10
        seqlen_q = seq.num_scheduled_tokens     # 本轮算多少
11
        end = start + seqlen_q
12
        input_ids.extend(seq[start:end])        # 只取要算的部分
13
        positions.extend(range(start, end))     # 位置编码
14
        # ... 累积长度用于 Flash Attention
15
        # ... 计算 slot_mapping 用于 KV-Cache 存储

1
def prepare_decode(self, seqs):
2
    for seq in seqs:
3
        input_ids.append(seq.last_token)        # 只有最后一个 token
4
        positions.append(len(seq) - 1)          # 最后的位置
5
        context_lens.append(len(seq))           # 当前总长度
6
        slot_mapping.append(                    # KV-Cache 写入位置
7
            seq.block_table[-1] * self.block_size + seq.last_block_num_tokens - 1
8
        )

1
def run(self, seqs, is_prefill):
2
    input_ids, positions = self.prepare_prefill(seqs) if is_prefill \
3
                           else self.prepare_decode(seqs)
4
    temperatures = self.prepare_sample(seqs)
5
    logits = self.run_model(input_ids, positions, is_prefill)
6
    token_ids = self.sampler(logits, temperatures).tolist()
7
    return token_ids

1
def capture_cudagraph(self):
2
    max_bs = min(max_num_seqs, 512)
3
    # 创建固定大小的输入张量
4
    input_ids = torch.zeros(max_bs, dtype=torch.int64)
5
    positions = torch.zeros(max_bs, dtype=torch.int64)
6
    # ...
7

8
    self.graph_bs = [1, 2, 4, 8, 16, 32, 48, 64, ..., 512]
9
    # 为每个批大小录制一个 CUDA Graph
10

11
    for bs in reversed(self.graph_bs):
12
        graph = torch.cuda.CUDAGraph()
13
        # 先跑一次让 CUDA 预热
14
        outputs[:bs] = self.model(input_ids[:bs], positions[:bs])
15
        # 录制
16
        with torch.cuda.graph(graph, self.graph_pool):
17
            outputs[:bs] = self.model(input_ids[:bs], positions[:bs])
18
        self.graphs[bs] = graph

1
def call(self, method_name, *args):
2
    if self.world_size > 1 and self.rank == 0:
3
        self.write_shm(method_name, *args)    # 主进程写共享内存
4
    method = getattr(self, method_name, None)
5
    return method(*args)

1
class LLMEngine:
2
    def __init__(self, model, **kwargs):
3
        # 1. 过滤出 Config 需要的参数
4
        config_fields = {field.name for field in fields(Config)}
5
        config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields}
6
        config = Config(model, **config_kwargs)
7

8
        # 2. 启动多 GPU 进程（如果需要）
9
        ctx = mp.get_context("spawn")
10
        for i in range(1, config.tensor_parallel_size):
11
            event = ctx.Event()
12
            process = ctx.Process(target=ModelRunner, args=(config, i, event))
13
            process.start()
14
            self.ps.append(process)
15
            self.events.append(event)
16

17
        # 3. 主进程也创建自己的 ModelRunner
18
        self.model_runner = ModelRunner(config, 0, self.events)
19

20
        # 4. 加载分词器
21
        self.tokenizer = AutoTokenizer.from_pretrained(config.model, use_fast=True)
22
        config.eos = self.tokenizer.eos_token_id
23

24
        # 5. 创建调度器
25
        self.scheduler = Scheduler(config)
26

27
        # 6. 注册退出回调
28
        atexit.register(self.exit)

1
def step(self):
2
    seqs, is_prefill = self.scheduler.schedule()                # 1. 调度：选谁算
3
    token_ids = self.model_runner.call("run", seqs, is_prefill) # 2. 执行：算
4
    self.scheduler.postprocess(seqs, token_ids, is_prefill)     # 3. 后处理：更新状态
5
    outputs = [(seq.seq_id, seq.completion_token_ids)           # 4. 收集完成的
6
               for seq in seqs if seq.is_finished]
7
    return outputs, num_tokens

1
def generate(self, prompts, sampling_params, use_tqdm=True):
2
    # 把所有 prompt 添加为请求
3
    for prompt, sp in zip(prompts, sampling_params):
4
        self.add_request(prompt, sp)
5

6
    # 循环直到所有请求完成
7
    while not self.is_finished():
8
        output, num_tokens = self.step()
9
        # num_tokens > 0 → prefill，=吞吐量
10
        # num_tokens < 0 → decode，绝对值是序列数
11
        for seq_id, token_ids in output:
12
            outputs[seq_id] = token_ids   # 收集完成的序列
13

14
    # 按顺序返回结果
15
    return [{"text": self.tokenizer.decode(token_ids), "token_ids": token_ids}
16
            for token_ids in outputs.values()]

1
class SiluAndMul(nn.Module):
2
    @torch.compile
3
    def forward(self, x):
4
        x, y = x.chunk(2, -1)     # 在最后一维切成两半
5
        return F.silu(x) * y      # 前一半过 SiLU，后一半直接乘

1
@triton.jit
2
def store_kvcache_kernel(key_ptr, key_stride, value_ptr, value_stride,
3
                          k_cache_ptr, v_cache_ptr, slot_mapping_ptr, D):
4
    idx = tl.program_id(0)
5
    slot = tl.load(slot_mapping_ptr + idx)
6
    if slot == -1: return                    # -1 表示不需要存储
7
    # 从 key/value 张量读取
8
    key = tl.load(key_ptr + idx * key_stride + tl.arange(0, D))
9
    value = tl.load(value_ptr + idx * value_stride + tl.arange(0, D))
10
    # 写入 KV-Cache 的指定位置
11
    tl.store(k_cache_ptr + slot * D + tl.arange(0, D), key)
12
    tl.store(v_cache_ptr + slot * D + tl.arange(0, D), value)

1
class Attention(nn.Module):
2
    def forward(self, q, k, v):
3
        context = get_context()
4
        # 第一步：把当前 K、V 存到缓存
5
        if k_cache.numel() and v_cache.numel():
6
            store_kvcache(k, v, k_cache, v_cache, context.slot_mapping)
7

8
        if context.is_prefill:
9
            if context.block_tables is not None:    # 有前缀缓存
10
                k, v = k_cache, v_cache             # 用缓存中的 K、V
11
            o = flash_attn_varlen_func(...)         # Flash Attention（变长序列）
12
        else:    # decode
13
            o = flash_attn_with_kvcache(...)        # Flash Attention（带缓存）
14
        return o

1
class QKVParallelLinear(ColumnParallelLinear):
2
    """专门为 QKV 投影设计的并行层"""
3
    def __init__(self, hidden_size, head_size, total_num_heads, total_num_kv_heads, bias):
4
        output_size = (total_num_heads + 2 * total_num_kv_heads) * head_size
5
        # 输出 = [Q区域 | K区域 | V区域]
6
        # Q: num_heads × head_size
7
        # K: num_kv_heads × head_size
8
        # V: num_kv_heads × head_size
9

10
    def weight_loader(self, param, loaded_weight, loaded_shard_id):
11
        # loaded_shard_id 是 "q"、"k" 或 "v"
12
        # 根据 shard_id 决定从 param 的哪个位置开始加载

1
class RMSNorm(nn.Module):
2
    def __init__(self, hidden_size, eps=1e-6):
3
        self.eps = eps
4
        self.weight = nn.Parameter(torch.ones(hidden_size))   # 可学习的缩放参数
5

6
    def rms_forward(self, x):
7
        x = x.float()
8
        var = x.pow(2).mean(dim=-1, keepdim=True)              # 计算均方值
9
        x.mul_(torch.rsqrt(var + self.eps))                    # 归一化：x / sqrt(var+eps)
10
        x = x.to(orig_dtype).mul_(self.weight)                 # 缩放
11
        return x

1
def add_rms_forward(self, x, residual):
2
    x = x.float().add_(residual.float())     # 先加残差
3
    residual = x.to(orig_dtype)              # 保存为下一层的残差
4
    # ... 然后做 RMSNorm

1
def apply_rotary_emb(x, cos, sin):
2
    x1, x2 = torch.chunk(x.float(), 2, dim=-1)    # 在最后一维切成两半
3
    y1 = x1 * cos - x2 * sin                      # 旋转变换（像2D旋转矩阵）
4
    y2 = x2 * cos + x1 * sin
5
    return torch.cat((y1, y2), dim=-1).to(x.dtype)

1
inv_freq = 1.0 / (base ** (torch.arange(0, rotary_dim, 2) / rotary_dim))

1
cache = torch.cat((cos, sin), dim=-1).unsqueeze_(1)
2
self.register_buffer("cos_sin_cache", cache, persistent=False)

1
@lru_cache(1)
2
def get_rope(head_size, rotary_dim, max_position, base):
3
    return RotaryEmbedding(head_size, rotary_dim, max_position, base)

1
class VocabParallelEmbedding(nn.Module):
2
    def __init__(self, num_embeddings, embedding_dim):
3
        self.tp_rank = dist.get_rank()        # 当前 GPU 编号
4
        self.tp_size = dist.get_world_size()  # 总 GPU 数
5
        # 把词表均匀分到多个 GPU 上
6
        self.num_embeddings_per_partition = num_embeddings // tp_size
7
        self.weight = nn.Parameter(torch.empty(num_embeddings_per_partition, embedding_dim))

1
def forward(self, x):
2
    if self.tp_size > 1:
3
        # 把不在本 GPU 范围的 token 置为 0
4
        mask = (x >= self.vocab_start_idx) & (x < self.vocab_end_idx)
5
        x = mask * (x - self.vocab_start_idx)     # 重新映射到本地索引
6
    y = F.embedding(x, self.weight)               # 查表
7
    if self.tp_size > 1:
8
        y = mask.unsqueeze(1) * y                 # 不属于本 GPU 的输出置 0
9
        dist.all_reduce(y)                        # 所有 GPU 求和
10
    return y

1
class ParallelLMHead(VocabParallelEmbedding):
2
    def forward(self, x):
3
        if context.is_prefill:
4
            # 只取每个序列的最后一个 token 的隐藏状态
5
            last_indices = context.cu_seqlens_q[1:] - 1
6
            x = x[last_indices].contiguous()
7
        logits = F.linear(x, self.weight)
8
        if self.tp_size > 1:
9
            # 每个 GPU 只有部分词表的 logits，需要收集（gather）到主 GPU
10
            dist.gather(logits, all_logits, 0)
11
            logits = torch.cat(all_logits, -1)
12
        return logits

1
class Sampler(nn.Module):
2
    @torch.compile
3
    def forward(self, logits, temperatures):
4
        logits = logits.float().div_(temperatures.unsqueeze(dim=1))  # 除以温度
5
        probs = torch.softmax(logits, dim=-1)                        # softmax 得到概率
6
        # Gumbel-Max 采样技巧
7
        sample_tokens = probs.div_(
8
            torch.empty_like(probs).exponential_(1).clamp_min_(1e-10)
9
        ).argmax(dim=-1)
10
        return sample_tokens

1
class Qwen3Attention(nn.Module):
2
    def forward(self, positions, hidden_states):
3
        qkv = self.qkv_proj(hidden_states)      # 1. QKV 投影（合并矩阵乘法）
4
        q, k, v = qkv.split([Q, K, V], -1)      # 2. 拆分为 Q、K、V
5
        q = q.view(-1, num_heads, head_dim)     # 3. 重塑为多头形状
6
        k = k.view(-1, num_kv_heads, head_dim)
7
        v = v.view(-1, num_kv_heads, head_dim)
8
        if not self.qkv_bias:
9
            q = self.q_norm(q)                  # 4. QK 归一化（Qwen3 特有）
10
            k = self.k_norm(k)
11
        q, k = self.rotary_emb(positions, q, k) # 5. 应用旋转位置编码
12
        o = self.attn(q, k, v)                  # 6. 注意力计算
13
        output = self.o_proj(o.flatten(1, -1))  # 7. 输出投影
14
        return output

1
class Qwen3MLP(nn.Module):
2
    def forward(self, x):
3
        gate_up = self.gate_up_proj(x)   # 合并的 Gate + Up 投影
4
        x = self.act_fn(gate_up)         # SiLU 门控激活
5
        x = self.down_proj(x)            # Down 投影
6
        return x

1
class Qwen3DecoderLayer(nn.Module):
2
    def forward(self, positions, hidden_states, residual):
3
        # 残差连接：hidden_states + residual
4
        if residual is None:
5
            hidden_states, residual = self.input_layernorm(hidden_states), hidden_states
6
        else:
7
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
8
        # 注意力
9
        hidden_states = self.self_attn(positions, hidden_states)
10
        # 残差 + 归一化
11
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
12
        # 前馈网络
13
        hidden_states = self.mlp(hidden_states)
14
        return hidden_states, residual

1
class Qwen3Model(nn.Module):
2
    def forward(self, input_ids, positions):
3
        hidden_states = self.embed_tokens(input_ids)    # 词嵌入
4
        residual = None
5
        for layer in self.layers:                        # 逐层计算
6
            hidden_states, residual = layer(positions, hidden_states, residual)
7
        hidden_states, _ = self.norm(hidden_states, residual)  # 最终归一化
8
        return hidden_states

1
class Qwen3ForCausalLM(nn.Module):
2
    packed_modules_mapping = {
3
        "q_proj": ("qkv_proj", "q"),
4
        "k_proj": ("qkv_proj", "k"),
5
        "v_proj": ("qkv_proj", "v"),
6
        "gate_proj": ("gate_up_proj", 0),
7
        "up_proj": ("gate_up_proj", 1),
8
    }

1
@dataclass(slots=True)
2
class Context:
3
    is_prefill: bool
4
    cu_seqlens_q: torch.Tensor | None       # 查询累积长度
5
    cu_seqlens_k: torch.Tensor | None       # 键累积长度
6
    max_seqlen_q: int                       # 最大查询长度
7
    max_seqlen_k: int                       # 最大键长度
8
    slot_mapping: torch.Tensor | None       # KV-Cache 槽位映射
9
    context_lens: torch.Tensor | None       # 各序列当前长度
10
    block_tables: torch.Tensor | None       # 块表

1
def load_model(model, path):
2
    packed_modules_mapping = getattr(model, "packed_modules_mapping", {})
3
    for file in glob(os.path.join(path, "*.safetensors")):
4
        with safe_open(file, "pt", "cpu") as f:
5
            for weight_name in f.keys():               # 遍历所有权重
6
                for k in packed_modules_mapping:        # 检查是否需要合并
7
                    if k in weight_name:
8
                        v, shard_id = packed_modules_mapping[k]
9
                        param_name = weight_name.replace(k, v)  # 映射到合并后的参数名
10
                        param = model.get_parameter(param_name)
11
                        weight_loader = getattr(param, "weight_loader")
12
                        weight_loader(param, f.get_tensor(weight_name), shard_id)
13
                        break
14
                else:
15
                    # 不需要合并，直接加载
16
                    param = model.get_parameter(weight_name)
17
                    weight_loader = getattr(param, "weight_loader", default_weight_loader)
18
                    weight_loader(param, f.get_tensor(weight_name))

1
path = "./Qwen3-0.6B/"
2
tokenizer = AutoTokenizer.from_pretrained(path)
3
llm = LLM(path, enforce_eager=True, tensor_parallel_size=1)
4

5
sampling_params = SamplingParams(temperature=0.6, max_tokens=256)
6
prompts = ["introduce yourself", "list all prime numbers within 100"]
7
prompts = [tokenizer.apply_chat_template(
8
    [{"role": "user", "content": prompt}],
9
    tokenize=False, add_generation_prompt=True,
10
) for prompt in prompts]
11

12
outputs = llm.generate(prompts, sampling_params)
13
for prompt, output in zip(prompts, outputs):
14
    print(f"Prompt: {prompt!r}")
15
    print(f"Completion: {output['text']!r}")