从零实现 LLM Inference:065. KV 元数据数组化(_block_infos / _block_refcounts)
KVBlockManager 的 block 元数据原来用 dict 做 global_id -> info/refcount 映射,decode 热路径会频繁查表。这里把两张表改成定长 list(按 global_id 直接索引),减少 Python dict 开销。
KVBlockManager 的 block 元数据原来用 dict 做 global_id -> info/refcount 映射,decode 热路径会频繁查表。这里把两张表改成定长 list(按 global_id 直接索引),减少 Python dict 开销。
KV append 的热路径里,每层每 token 都要更新一次 block length。之前用 NamedTuple 需要不断创建新对象并回写 dict;这版改成 slots dataclass,length 原地自增,减少 Python 分配和重复查表。
prefix cache hit 后,多条 session 会共享同一批 KV blocks;第一次 decode 写入时如果 last block 还没写满,就会触发 COW:先 clone block 再 append token。原实现用 index_select/index_copy 搬整块 KV;这版加...
上一版 full-batch KV-append 只覆盖 pos 常量的稳态;但只要 prompt 长度不一致,decode 的 pos 就会按 request 分叉,仍然要构造 batch_idx/pos 这类 index tensor。这版补一个 identity batch 的 Triton kernel:...
batch decode 稳态里,经常满足 fast_batch_idx 是 [0..B-1] 且同一步 pos 对整个 batch 是常量;这时 Triton KV append 还在每 step 构造 batch_idx/pos 这类小 tensor。这个 PR 新增 full-batch Triton ke...
append_token_batch 里 fast_batch_idx 很多时候就是 [0..B-1];之前每层都会 index_select 把 key_new/value_new 重新拷一遍,还会构造 pos_t。这个小改动在 identity batch 时直接复用 key_new/value_new,并且...
append_token_batch 之前只覆盖 len<block_size 的 fast path;一旦 last block 满了(len==block_size)就会退化成逐 request 的 append_token。这个点会制造 ITL 的尖刺。这版把 rollover 也塞回 batch:先...
prefix cache hit 时我们还在把 last_logits 从 CPU 拷回 GPU;这版把 entry 的 last_logits 直接存成 device 上的一份小 clone,hit 变成真正的零拷贝。
prefix cache 复用会让 decode 的最后一个 KV block refcount>1;之前 append_token_batch 直接退化成逐 request 的 append_token + copy-on-write,CPU/GPU overhead 都很明显。这版把 COW clone...
055 做了 longest-prefix reuse,但 longest-prefix 查询还是 O(N) 扫描;这版用 token trie 替换掉 scan,把 cache miss 的 longest-prefix 查找从 ms 级降到 us 级,减少 scheduler CPU overhead。
prefix cache 之前只能 exact hit;这版做 longest-prefix 复用:命中“缓存 prompt 是新 prompt 的前缀”时直接挂载 KV blocks,然后用 decode(T=1) teacher-forcing 补齐 suffix。顺手把 paged-attn Triton ...
prefix cache 之前用 prompt 字符串当 key;在 pretok 场景里,prompt 文本不同但 token ids 相同会导致 cache miss。改成优先用 prompt_token_ids tuple 作为 key,并加了一个 benchmark knob 复现/量化收益。
paged-attn + CUDA Graph decode 里,每步都在 torch.tensor(list) 然后 copy_ 到 GPU,既有分配也不是真正 non_blocking。改成复用 pinned host buffer + numpy view 直接写入,吞吐提升、TPOT 下降。
SchedulerManager 的 worker loop 里每个 token 都会进一次 lock 取 q/detok/state;改成每步 decode 批量抓取一次,减少 lock acquire/release 和热点 dict 访问。
之前 top_k+top_p 采样每步都会对整个 vocab 做 sort/gather,开销巨大。把采样改成直接在 sorted topk 空间里做 top-p + multinomial,并顺手把 top_k clamp 到 vocab size。
streaming 如果每个 token 都 q.put + HTTP flush,会把 Python/IO overhead 放大。引入 stream_interval:第一段内容立刻发,后续每 N 个 token 合并成一个 chunk 再发,吞吐和 tail latency 都能更稳。
burst streaming 下 unlimited inflight 会把排队延迟炸到 p99,但吞吐几乎不变。给 SchedulerManager 加 max_inflight_requests:超过就直接拒绝(HTTP 429),benchmark_streaming 也打印 Rejected 与 rej...
submit_interval_ms 如果用“sleep after each add_request”的相对口径,提交路径一变快,实际到达率就变了,TTFT 会被排队形态污染。给 benchmark_streaming 加 submit_scheudle=absolute,用 t0+i*interval 固定到...
add_request 里同步做 tokenizer.encode 会放大 submit wall / p99。引入 tokenize_workers 线程池,把 tokenization 变成后台阶段,并把 TTFT 拆到 tokenize/queue/prefill 三段。
streaming bench 里 add_request 的 p99 很容易被 tokenizer.encode 的 CPU 开销污染。增加 add_request(prompt_token_ids=…) + benchmark_streaming –pretok,把 tokenization 从提交路径挪出去。
TTFT 只是一个总数:里面既有排队等待,也有 prefill 的真实计算。只看 TTFT 很容易把优化方向搞反。给 SchedulerManager 记录 admit timestamp,并在 benchmark_streaming 里打印 TTFT breakdown。
paged-attn / CUDA Graph 的第一次请求会把 Triton JIT 和 graph capture 算进 TTFT/吞吐,导致对比结论失真。给 benchmark_streaming 加 warmup-runs + repeat-runs,把冷启动和稳态拆开。
streaming 场景想看 TTFT/TPOT/ITL 的 p99,但之前不好一条命令切换 decode 的快路径,也不好给 nsys 挂 NVTX。补齐 benchmark 的三个开关。
decode 每层都要把 (k,v) 写进 KV cache。原来是 index_select + index_put 四连。先用一个 Triton kernel fuse 成一次写入,并加上 batch size gate,避免小 batch 退化。
decode 已经很快了但还不够?很多时候慢在 sampling:每步 B 次 .item() 会把 GPU pipeline 打散。把采样做成 batch,一步只同步一次。
decode 热路径里 kernel launch 太密?把一次 decode step 捕获成 CUDA Graph,replay 省掉大量 CPU dispatch。
把 paged-attn decode kernel 的 num_warps/num_stages 从“拍脑袋常量”变成 Triton autotune;结果在 decode 热路径里反而回退了。
paged decode 还在慢?很多时候瓶颈不在 attention,而是在每步 L*B 次的 KV 写入:做一个 batch fast-path,直接把 Python 循环砍掉。
block_table 不再每步 gather/copy:把它常驻 GPU,用 slot_mapping 让 triton kernel 直接索引,decode 吞吐更稳。
paged attention decode 每步重建 block_table 太浪费:按 session 缓存 row,只在 block 变化时更新,decode 吞吐和 TPOT 更稳。
paged attention decode 路径里不该构建 attention_mask;同时把 block_tables 的 H2D copy 合并成一次 async,收一点点 TPOT/throughput。
pack admission 会把 active sessions 拉得很高:加一个 max_active_requests(max_num_seqs)把 decode backlog 的 ITL/TPOT tail 收回来。
token budget + FIFO 仍然会遇到 head-of-line blocking:用 lookahead packing 把短请求先塞进 prefill,收敛 TTFT p99。
worker loop 在有 active sessions 时先 decode 再 prefill:减少 decode 被 admission 抢占,收敛 ITL p99。
prefill admission 从“按请求数”升级为“按 tokens 预算”:限制 prefill 抢占,收敛 ITL p99。
把 worker 的 prefill admission batch size 从 decode batch size 解耦:TTFT p99 直接砍掉一半。
prefix cache 开启时也能 micro-batch prefill:hit 直接 attach,miss 合并成一次 forward。
把 worker 里的 prefill 从串行变成 micro-batch,一次 forward 吞掉一批 pending request。
把 streaming 的 add_request 从 prefill 中解耦:快速入队,由 worker 统一做 prefill + decode。
OnlineScheduler 支持 prompt_token_ids,server 去掉重复 encode,benchmark 增加 –pretok。
把 server/benchmark 对 OnlineScheduler._sessions 的直接访问收口。
补齐 SchedulerManager 的 close,优雅停掉 worker 并结束 streaming。
server worker 用 threading.Event 驱动唤醒,去掉 idle 轮询 sleep。
OnlineScheduler 用 deque 维护活跃队列,去掉每 step 的全量扫描。
给 benchmark 加上 TTFT/TPOT,并把 server 的 finished 清理从扫描改成事件。
支持从 HuggingFace 加载 GPT2 权重,为后续和 vLLM/sglang 对齐 benchmark 铺路。
实现真正的 paged attention。
通过性能观测进行性能优化。
使用 pytorch profiler 进行性能观测。
实现简单的 prefix caching,通过 prefix cache 来复用之前的 kv-cache。
实现简单的 benchmark,对比不同实现的性能。
实现 scheduler manager,支持 online scheduler 的接入。
支持简单的 openai api,实现 chat completion。
实现简单的 inference server,使用 FastAPI 以及 uvicorn。
实现 online scheduler,展示连续批处理。
让 kv block manager 真正发挥作用,实现 python 版 paged attention。
实现 batch decode,并隐式实现 continuous batching。
实现基础的 kv block manager
实现离线调度器,支持并发请求。
添加 Inference Session,支持并发请求。
实现流式生成(streaming),支持边生成边输出。
解决模型初始化问题,让训练更稳定。
实现 batch inference 功能,支持多条 prompt 同时推理。
实现 top-k top-p 这种 sampling 操作,并把整体的 prefill decode 流程规范化,对齐业界 vllm,huggingface 的实现。
实现 kv-cache 部分,让模型能够处理推理时的 kv-cache。
实现最基本的 greedy generate。
添加单独的 eval.py 文件,从而将评估与训练相分离。
实现梯度累积和梯度裁剪,提高训练稳定性。
引入 FineWebNPYDataset 数据集 class。
使用 WandB 记录训练过程,方便后续分析。
使用 PyTorch profiler 与 NVTX 捕捉 trace,深入分析训练性能瓶颈。
通过 activation checkpointing 以重计算换显存,优化大模型训练。
为学习率引入 cosine scheduler,并将调度状态写入 checkpoint。
将模型配置写入 checkpoint,简化加载与推理时的配置管理。
重构 Dataset,支持多文档输入与更合理的切分策略。
在完成基础训练后,实现一个最简单的文本生成脚本。
将 toy 数据集替换为更真实的文本数据,完善训练链路。
加入验证集评估和日志记录,用 loss 与 PPL 监控训练效果。
为训练脚本引入 argparse 命令行参数,向工业级实现迈进。
为训练过程加入 checkpoint 容错机制,支持从中间状态恢复。
在完成张量并行后,引入混合精度训练以提高算力利用率。
通过按 head 维度切分 QKV,让 Attention 形成真正的张量并行。
将 Row Parallel Linear 引入 Attention 层,搭配 Column Parallel 完成张量并行。
把 Row Parallel Linear 引入 FFN 模块,配合 Column Parallel 完成张量并行。
在已有 Column Parallel 基础上实现 Row Parallel Linear。
把 Column Parallel Linear 集成到 GPTModel 中并用训练循环验证。
引入列张量并行(Column Parallel)作为张量并行的第一步。
使用 PyTorch DDP 实现最简单的数据并行训练。
为 mini-GPT 加上最基本的 loss 和最小 train loop。
从零实现 mini-GPT 的配置、模型和简单前向。
Generalized Advantage Estimation(GAE)推导与直觉
数据并行
训练时的显存开销
SGLang
KVBlockManager 的 block 元数据原来用 dict 做 global_id -> info/refcount 映射,decode 热路径会频繁查表。这里把两张表改成定长 list(按 global_id 直接索引),减少 Python dict 开销。
KV append 的热路径里,每层每 token 都要更新一次 block length。之前用 NamedTuple 需要不断创建新对象并回写 dict;这版改成 slots dataclass,length 原地自增,减少 Python 分配和重复查表。
prefix cache hit 后,多条 session 会共享同一批 KV blocks;第一次 decode 写入时如果 last block 还没写满,就会触发 COW:先 clone block 再 append token。原实现用 index_select/index_copy 搬整块 KV;这版加...
上一版 full-batch KV-append 只覆盖 pos 常量的稳态;但只要 prompt 长度不一致,decode 的 pos 就会按 request 分叉,仍然要构造 batch_idx/pos 这类 index tensor。这版补一个 identity batch 的 Triton kernel:...
batch decode 稳态里,经常满足 fast_batch_idx 是 [0..B-1] 且同一步 pos 对整个 batch 是常量;这时 Triton KV append 还在每 step 构造 batch_idx/pos 这类小 tensor。这个 PR 新增 full-batch Triton ke...
append_token_batch 里 fast_batch_idx 很多时候就是 [0..B-1];之前每层都会 index_select 把 key_new/value_new 重新拷一遍,还会构造 pos_t。这个小改动在 identity batch 时直接复用 key_new/value_new,并且...
append_token_batch 之前只覆盖 len<block_size 的 fast path;一旦 last block 满了(len==block_size)就会退化成逐 request 的 append_token。这个点会制造 ITL 的尖刺。这版把 rollover 也塞回 batch:先...
prefix cache hit 时我们还在把 last_logits 从 CPU 拷回 GPU;这版把 entry 的 last_logits 直接存成 device 上的一份小 clone,hit 变成真正的零拷贝。
prefix cache 复用会让 decode 的最后一个 KV block refcount>1;之前 append_token_batch 直接退化成逐 request 的 append_token + copy-on-write,CPU/GPU overhead 都很明显。这版把 COW clone...
055 做了 longest-prefix reuse,但 longest-prefix 查询还是 O(N) 扫描;这版用 token trie 替换掉 scan,把 cache miss 的 longest-prefix 查找从 ms 级降到 us 级,减少 scheduler CPU overhead。
prefix cache 之前只能 exact hit;这版做 longest-prefix 复用:命中“缓存 prompt 是新 prompt 的前缀”时直接挂载 KV blocks,然后用 decode(T=1) teacher-forcing 补齐 suffix。顺手把 paged-attn Triton ...
prefix cache 之前用 prompt 字符串当 key;在 pretok 场景里,prompt 文本不同但 token ids 相同会导致 cache miss。改成优先用 prompt_token_ids tuple 作为 key,并加了一个 benchmark knob 复现/量化收益。
paged-attn + CUDA Graph decode 里,每步都在 torch.tensor(list) 然后 copy_ 到 GPU,既有分配也不是真正 non_blocking。改成复用 pinned host buffer + numpy view 直接写入,吞吐提升、TPOT 下降。
SchedulerManager 的 worker loop 里每个 token 都会进一次 lock 取 q/detok/state;改成每步 decode 批量抓取一次,减少 lock acquire/release 和热点 dict 访问。
之前 top_k+top_p 采样每步都会对整个 vocab 做 sort/gather,开销巨大。把采样改成直接在 sorted topk 空间里做 top-p + multinomial,并顺手把 top_k clamp 到 vocab size。
streaming 如果每个 token 都 q.put + HTTP flush,会把 Python/IO overhead 放大。引入 stream_interval:第一段内容立刻发,后续每 N 个 token 合并成一个 chunk 再发,吞吐和 tail latency 都能更稳。
burst streaming 下 unlimited inflight 会把排队延迟炸到 p99,但吞吐几乎不变。给 SchedulerManager 加 max_inflight_requests:超过就直接拒绝(HTTP 429),benchmark_streaming 也打印 Rejected 与 rej...
submit_interval_ms 如果用“sleep after each add_request”的相对口径,提交路径一变快,实际到达率就变了,TTFT 会被排队形态污染。给 benchmark_streaming 加 submit_scheudle=absolute,用 t0+i*interval 固定到...
add_request 里同步做 tokenizer.encode 会放大 submit wall / p99。引入 tokenize_workers 线程池,把 tokenization 变成后台阶段,并把 TTFT 拆到 tokenize/queue/prefill 三段。
streaming bench 里 add_request 的 p99 很容易被 tokenizer.encode 的 CPU 开销污染。增加 add_request(prompt_token_ids=…) + benchmark_streaming –pretok,把 tokenization 从提交路径挪出去。
TTFT 只是一个总数:里面既有排队等待,也有 prefill 的真实计算。只看 TTFT 很容易把优化方向搞反。给 SchedulerManager 记录 admit timestamp,并在 benchmark_streaming 里打印 TTFT breakdown。
paged-attn / CUDA Graph 的第一次请求会把 Triton JIT 和 graph capture 算进 TTFT/吞吐,导致对比结论失真。给 benchmark_streaming 加 warmup-runs + repeat-runs,把冷启动和稳态拆开。
streaming 场景想看 TTFT/TPOT/ITL 的 p99,但之前不好一条命令切换 decode 的快路径,也不好给 nsys 挂 NVTX。补齐 benchmark 的三个开关。
decode 每层都要把 (k,v) 写进 KV cache。原来是 index_select + index_put 四连。先用一个 Triton kernel fuse 成一次写入,并加上 batch size gate,避免小 batch 退化。
decode 已经很快了但还不够?很多时候慢在 sampling:每步 B 次 .item() 会把 GPU pipeline 打散。把采样做成 batch,一步只同步一次。
decode 热路径里 kernel launch 太密?把一次 decode step 捕获成 CUDA Graph,replay 省掉大量 CPU dispatch。
把 paged-attn decode kernel 的 num_warps/num_stages 从“拍脑袋常量”变成 Triton autotune;结果在 decode 热路径里反而回退了。
paged decode 还在慢?很多时候瓶颈不在 attention,而是在每步 L*B 次的 KV 写入:做一个 batch fast-path,直接把 Python 循环砍掉。
block_table 不再每步 gather/copy:把它常驻 GPU,用 slot_mapping 让 triton kernel 直接索引,decode 吞吐更稳。
paged attention decode 每步重建 block_table 太浪费:按 session 缓存 row,只在 block 变化时更新,decode 吞吐和 TPOT 更稳。
paged attention decode 路径里不该构建 attention_mask;同时把 block_tables 的 H2D copy 合并成一次 async,收一点点 TPOT/throughput。
pack admission 会把 active sessions 拉得很高:加一个 max_active_requests(max_num_seqs)把 decode backlog 的 ITL/TPOT tail 收回来。
token budget + FIFO 仍然会遇到 head-of-line blocking:用 lookahead packing 把短请求先塞进 prefill,收敛 TTFT p99。
worker loop 在有 active sessions 时先 decode 再 prefill:减少 decode 被 admission 抢占,收敛 ITL p99。
prefill admission 从“按请求数”升级为“按 tokens 预算”:限制 prefill 抢占,收敛 ITL p99。
把 worker 的 prefill admission batch size 从 decode batch size 解耦:TTFT p99 直接砍掉一半。
prefix cache 开启时也能 micro-batch prefill:hit 直接 attach,miss 合并成一次 forward。
把 worker 里的 prefill 从串行变成 micro-batch,一次 forward 吞掉一批 pending request。
把 streaming 的 add_request 从 prefill 中解耦:快速入队,由 worker 统一做 prefill + decode。
OnlineScheduler 支持 prompt_token_ids,server 去掉重复 encode,benchmark 增加 –pretok。
把 server/benchmark 对 OnlineScheduler._sessions 的直接访问收口。
补齐 SchedulerManager 的 close,优雅停掉 worker 并结束 streaming。
server worker 用 threading.Event 驱动唤醒,去掉 idle 轮询 sleep。
OnlineScheduler 用 deque 维护活跃队列,去掉每 step 的全量扫描。
给 benchmark 加上 TTFT/TPOT,并把 server 的 finished 清理从扫描改成事件。
支持从 HuggingFace 加载 GPT2 权重,为后续和 vLLM/sglang 对齐 benchmark 铺路。
实现真正的 paged attention。
通过性能观测进行性能优化。
使用 pytorch profiler 进行性能观测。
实现简单的 prefix caching,通过 prefix cache 来复用之前的 kv-cache。
实现简单的 benchmark,对比不同实现的性能。
实现 scheduler manager,支持 online scheduler 的接入。
支持简单的 openai api,实现 chat completion。
实现简单的 inference server,使用 FastAPI 以及 uvicorn。
实现 online scheduler,展示连续批处理。
让 kv block manager 真正发挥作用,实现 python 版 paged attention。
实现 batch decode,并隐式实现 continuous batching。
实现基础的 kv block manager
实现离线调度器,支持并发请求。
添加 Inference Session,支持并发请求。
实现流式生成(streaming),支持边生成边输出。
实现 batch inference 功能,支持多条 prompt 同时推理。
实现 top-k top-p 这种 sampling 操作,并把整体的 prefill decode 流程规范化,对齐业界 vllm,huggingface 的实现。
实现 kv-cache 部分,让模型能够处理推理时的 kv-cache。
实现最基本的 greedy generate。
解决模型初始化问题,让训练更稳定。
添加单独的 eval.py 文件,从而将评估与训练相分离。
实现梯度累积和梯度裁剪,提高训练稳定性。
引入 FineWebNPYDataset 数据集 class。
使用 WandB 记录训练过程,方便后续分析。
使用 PyTorch profiler 与 NVTX 捕捉 trace,深入分析训练性能瓶颈。
通过 activation checkpointing 以重计算换显存,优化大模型训练。
为学习率引入 cosine scheduler,并将调度状态写入 checkpoint。
将模型配置写入 checkpoint,简化加载与推理时的配置管理。
重构 Dataset,支持多文档输入与更合理的切分策略。
在完成基础训练后,实现一个最简单的文本生成脚本。
将 toy 数据集替换为更真实的文本数据,完善训练链路。
加入验证集评估和日志记录,用 loss 与 PPL 监控训练效果。
为训练脚本引入 argparse 命令行参数,向工业级实现迈进。
为训练过程加入 checkpoint 容错机制,支持从中间状态恢复。
在完成张量并行后,引入混合精度训练以提高算力利用率。
通过按 head 维度切分 QKV,让 Attention 形成真正的张量并行。
将 Row Parallel Linear 引入 Attention 层,搭配 Column Parallel 完成张量并行。
把 Row Parallel Linear 引入 FFN 模块,配合 Column Parallel 完成张量并行。
在已有 Column Parallel 基础上实现 Row Parallel Linear。
把 Column Parallel Linear 集成到 GPTModel 中并用训练循环验证。
引入列张量并行(Column Parallel)作为张量并行的第一步。
使用 PyTorch DDP 实现最简单的数据并行训练。
为 mini-GPT 加上最基本的 loss 和最小 train loop。
从零实现 mini-GPT 的配置、模型和简单前向。
回归基础,返璞归真
Collections of leetcode contests.
Collections of codeforces contests.
Collections of high frequent problems.
Collections of algorithm template.
Practices of some codeforces constructive problems.
交叉熵中的数学
交叉熵中的数学