Compare commits
22
Commits
b2ff243447
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
 ViperEkura
|
9de0bad3d4 | ||
|
ViperEkura
|
d471cfa276 | ||
|
ViperEkura
|
12d587aa92 | ||
|
ViperEkura
|
4f14d09fe3 | ||
|
ViperEkura
|
bb0c32b032 | ||
|
ViperEkura
|
c9f290c3c8 | ||
|
ViperEkura
|
f5499866ad | ||
|
ViperEkura
|
740e766e99 | ||
|
ViperEkura
|
01ccf192ea | ||
|
ViperEkura
|
c2c18f10d6 | ||
|
ViperEkura
|
e594d02d31 | ||
|
ViperEkura
|
29007c7b97 | ||
|
ViperEkura
|
6b0a1dbb5e | ||
|
ViperEkura
|
c05a432e45 | ||
|
ViperEkura
|
57abefa47f | ||
|
ViperEkura
|
6b26ec33ab | ||
|
ViperEkura
|
0018868ee3 | ||
|
ViperEkura
|
496f964979 | ||
|
ViperEkura
|
4d96a84fc5 | ||
|
ViperEkura
|
e7d736a3b0 | ||
|
ViperEkura
|
eeaf0a5a16 | ||
|
ViperEkura
|
ba100c19d1 |
+12
-10
@@ -1,10 +1,12 @@
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.pyo
|
||||
.venv/
|
||||
venv/
|
||||
*.egg-info/
|
||||
dist/
|
||||
build/
|
||||
output/
|
||||
.DS_Store
|
||||
# Ignore everything by default
|
||||
*
|
||||
|
||||
# Recurse into directories
|
||||
!*/
|
||||
|
||||
|
||||
# Whitelist: only source code and docs
|
||||
!*.py
|
||||
!README.md
|
||||
!AGENTS.md
|
||||
!.gitignore
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
# AGENTS.md — AstrAI Promo
|
||||
|
||||
## Project
|
||||
|
||||
Manim-driven technical animation project for AstrAI promo videos (1B param, zh/en bilingual LLM).
|
||||
Python 3.12, Manim (renders MP4 output).
|
||||
|
||||
## Setup
|
||||
|
||||
Requires a conda environment named `manim`.
|
||||
|
||||
```bash
|
||||
conda run -n manim python -m pip install manim
|
||||
```
|
||||
|
||||
The user verbally referred to it as "animin" — the actual name is `manim`. No system Python is available; always use `conda run -n manim`.
|
||||
|
||||
## Commands
|
||||
|
||||
| Action | Command |
|
||||
|--------|---------|
|
||||
| Render all (parallel) | `conda run -n manim python render_all.py` |
|
||||
| Render all (draft 480p) | `conda run -n manim python render_all.py -ql` |
|
||||
| Render all (4K) | `conda run -n manim python render_all.py -4k` |
|
||||
| Single scene (1080p) | `conda run -n manim python -m manim <file.py> <SceneClass> -qh` |
|
||||
| Single scene (draft) | `conda run -n manim python -m manim <file.py> <SceneClass> -ql` |
|
||||
|
||||
Output goes to `output/videos/`.
|
||||
|
||||
## File layout (flat — no subdirectory)
|
||||
|
||||
Files are at the repo root, **not** under `promo/` despite what `README.md` and `PROMO_GUIDE.md` reference in some spots.
|
||||
|
||||
```
|
||||
hook.py # Scene: Hook
|
||||
architecture.py # Scene: Architecture
|
||||
continuous_batching.py # Scene: ContinuousBatching
|
||||
cta.py # Scene: CTA
|
||||
paged_cache.py # Scene: PrefixCache
|
||||
transformer.py # Scene: Transformer
|
||||
render_all.py # Parallel batch renderer
|
||||
PROMO_GUIDE.md # Full production guide (storyboard, narration, recording)
|
||||
README.md # Quick-start docs (may reference stale promo/ paths)
|
||||
```
|
||||
|
||||
## Scenes
|
||||
|
||||
| File | Class | Content | Est. duration |
|
||||
|------|-------|---------|---------------|
|
||||
| `hook.py` | `Hook` | Terminal demo + GPU comparison + logo reveal | ~18s |
|
||||
| `transformer.py` | `Transformer` | GQA attention (Q/K/V → RoPE → Attention → O) + spec card | ~35s |
|
||||
| `continuous_batching.py` | `ContinuousBatching` | 4-stage pipeline + throughput comparison (1.0x vs 3.4x) | ~30s |
|
||||
| `paged_cache.py` | `PrefixCache` | Paged KV cache: page table, on-demand growth, O(1) free | ~30s |
|
||||
| `architecture.py` | `Architecture` | Full-stack: HTTP API → Engine → Paged Cache → Transformer | ~25s |
|
||||
| `cta.py` | `CTA` | End card: "AstrAI — Single GPU, Open Source, 1B params" | ~3s |
|
||||
|
||||
## Render quirks
|
||||
|
||||
- `render_all.py` runs all 5 scenes **in parallel** via `ThreadPoolExecutor`. Script blocks until all finish. No `-qh` flag in render_all.py means check the default at the top of the file.
|
||||
- The `output/` directory must exist (script creates it via Manim's `--media_dir`).
|
||||
|
||||
## Chinese text in Manim
|
||||
|
||||
```python
|
||||
Text.set_default(font="Microsoft YaHei")
|
||||
```
|
||||
This must be set before creating any `Text` objects, otherwise CJK characters render as boxes.
|
||||
|
||||
## No tests / lint / typecheck
|
||||
|
||||
This is a render-only project. No pytest, ruff, mypy, or similar config exists.
|
||||
|
||||
## .gitignore oddity (from parent AstrAI project)
|
||||
|
||||
The repo uses deny-by-default `.gitignore` (`*` then `!*.py`, `!*.sh`, etc.).
|
||||
New file types at root need an explicit `!` entry — check `.gitignore` before adding assets.
|
||||
-373
@@ -1,373 +0,0 @@
|
||||
# AstrAI 宣传视频制作指南
|
||||
|
||||
> 本文档为制作 AstrAI 宣传视频提供完整的技术参考、分镜建议和录制脚本。
|
||||
> 目标时长：**2-3 分钟**
|
||||
|
||||
---
|
||||
|
||||
## 目录
|
||||
|
||||
1. [项目定位与核心卖点](#1-项目定位与核心卖点)
|
||||
2. [技术架构速览](#2-技术架构速览)
|
||||
3. [分镜脚本](#3-分镜脚本)
|
||||
4. [演示录制指南](#4-演示录制指南)
|
||||
5. [动画场景说明](#5-动画场景说明)
|
||||
6. [旁白文案草稿](#6-旁白文案草稿)
|
||||
7. [素材清单](#7-素材清单)
|
||||
|
||||
---
|
||||
|
||||
## 1. 项目定位与核心卖点
|
||||
|
||||
**一句话定位：**
|
||||
> 一个能在单张消费级 GPU 上训练和推理的 1B 参数中英双语语言模型框架。
|
||||
|
||||
**核心卖点（视频中需突出）：**
|
||||
|
||||
| 卖点 | 说明 | 视觉表达 |
|
||||
|------|------|---------|
|
||||
| **单卡可跑** | 1B 参数，RTX 3090/4090 即可运行 | 巨大服务器集群 vs 单张显卡对比 |
|
||||
| **连续批处理** | 动态合并请求，吞吐量 3x+ | 任务流经 Cleanup→Refill→Prefill→Decode 动画 |
|
||||
| **前缀缓存零拷贝** | 相同前缀直接复用 KV，无需重算 | Radix Tree 生长动画 |
|
||||
| **OpenAI 兼容 API** | 一行代码切换 | curl 命令对比 |
|
||||
| **流式输出** | 逐 token 返回，低首延迟 | 终端逐字喷出效果 |
|
||||
| **全过程开源** | 训练+推理+权重全部开源 | GitHub 页面展示 |
|
||||
|
||||
---
|
||||
|
||||
## 2. 技术架构速览
|
||||
|
||||
### 整体架构
|
||||
|
||||
```
|
||||
┌──────────────────────────────────────────────────┐
|
||||
│ FastAPI Server (OpenAI-Compatible API) │
|
||||
├──────────────────────────────────────────────────┤
|
||||
│ InferenceEngine (Streaming + Async + Batch) │
|
||||
├──────────────────────────────────────────────────┤
|
||||
│ Continuous Batching Scheduler │
|
||||
│ ┌────────┐ ┌──────┐ ┌────────┐ ┌────────┐ │
|
||||
│ │Cleanup │→ │Refill│→ │Prefill │→ │ Decode │ │
|
||||
│ └────────┘ └──────┘ └────────┘ └────────┘ │
|
||||
├──────────────────────────────────────────────────┤
|
||||
│ Prefix Cache (Radix Tree) + KV Cache │
|
||||
├──────────────────────────────────────────────────┤
|
||||
│ Transformer (24层 GQA, RoPE, SwiGLU) │
|
||||
└──────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### 关键技术指标
|
||||
|
||||
| 指标 | 值 |
|
||||
|------|------|
|
||||
| 参数量 | ~1.0B |
|
||||
| 词表大小 | 100,000（中英 BPE） |
|
||||
| 层数 | 24 |
|
||||
| 注意力头 | 24 Q-heads / 4 KV-heads（GQA） |
|
||||
| 最大长度 | 2048 tokens |
|
||||
| 精度 | bfloat16 |
|
||||
| 最低显存 | ~6GB（推理）/~12GB（训练） |
|
||||
|
||||
---
|
||||
|
||||
## 3. 分镜脚本
|
||||
|
||||
总时长 **2:30**，分为 6 个段落。
|
||||
|
||||
### Segment 1：Hook + 问题陈述（0:00 - 0:20）
|
||||
|
||||
| 镜头 | 画面 | 旁白 | 时长 |
|
||||
|------|------|------|------|
|
||||
| 1.1 | 黑屏，逐字打出"大语言模型很强大" | "大语言模型很强大——" | 3s |
|
||||
| 1.2 | 切到数据中心照片 / 巨大 GPU 集群 | "——但跑起来需要几十张 GPU，普通人根本碰不到。" | 5s |
|
||||
| 1.3 | 画面分屏：左边集群，右边一张 RTX 4090 | "但如果我告诉你，只要一张显卡就够了呢？" | 5s |
|
||||
| 1.4 | Logo 出现：**AstrAI**，下方副标题 "1B 参数单卡推理框架" | "AstrAI——单卡跑大模型。" | 7s |
|
||||
|
||||
**视觉素材**：数据中心图片（可免版权下载）、RTX 4090 产品图、Logo 动画
|
||||
|
||||
---
|
||||
|
||||
### Segment 2：模型架构速览（0:20 - 0:45）
|
||||
|
||||
| 镜头 | 画面 | 旁白 | 时长 |
|
||||
|------|------|------|------|
|
||||
| 2.1 | Transformer 架构图逐层展开：embed → 24层 decoder → norm → head | "1B 参数，24 层 Transformer，100k 词表的中英 BPE 分词器。" | 8s |
|
||||
| 2.2 | 高亮 GQA：24个 Q head 映射到 4个 KV head | "GQA 分组查询注意力——24 个查询头只对应 4 个 KV 头，KV 缓存直接减少 83%。" | 10s |
|
||||
| 2.3 | RoPE 旋转变换可视化 | "RoPE 旋转位置编码，支持动态长度外推。" | 5s |
|
||||
| 2.4 | fade 到模型 card：vocab=100k, dim=1536, layers=24, heads=24, kv_heads=4 | 静默 | 2s |
|
||||
|
||||
**视觉素材**：`architecture.py` 动画、模型参数 card
|
||||
|
||||
---
|
||||
|
||||
### Segment 3：连续批处理（0:45 - 1:20）
|
||||
|
||||
| 镜头 | 画面 | 旁白 | 时长 |
|
||||
|------|------|------|------|
|
||||
| 3.1 | 3 个请求同时到达服务器 | "当多个请求同时到达时——" | 3s |
|
||||
| 3.2 | 静态批处理对比：最长补齐，3个请求串行 → 总耗时 max_len × 3 | "传统做法是静态批处理，把请求补齐到相同长度，串行处理，GPU 利用率低下。" | 8s |
|
||||
| 3.3 | 连续批处理动画：任务流入 Waiting Queue → Cleanup → Refill → Prefill → Decode | "AstrAI 采用连续批处理：任务动态进出，GPU 每一刻都在满负荷运转。" | 10s |
|
||||
| 3.4 | 放大 Decode 阶段：同一位置的任务合并成一批 | "特别地，只有处于相同 KV 缓存位置的任务才一起解码，从根本上避免了 RoPE 位置错乱的问题。" | 8s |
|
||||
| 3.5 | 吞吐对比柱状图：Static Batch vs Continuous Batching (3x+) | "实测吞吐量提升 3 倍以上。" | 6s |
|
||||
|
||||
**视觉素材**：`continuous_batching.py` 动画、对比图表
|
||||
|
||||
---
|
||||
|
||||
### Segment 4：前缀缓存（1:20 - 1:50）
|
||||
|
||||
| 镜头 | 画面 | 旁白 | 时长 |
|
||||
|------|------|------|------|
|
||||
| 4.1 | 两个请求有相同 system prompt："你是一个AI助手" | "如果两个请求有相同的前缀——比如相同的系统提示词——" | 5s |
|
||||
| 4.2 | 普通做法：两个请求各自独立计算前 20 个 token | "普通框架会各自从头计算一遍，白白浪费算力。" | 5s |
|
||||
| 4.3 | Radix Tree 生长动画：第一个请求插入，第二个请求匹配共享前缀 | "AstrAI 用一颗字典树缓存所有前缀的 KV——第二个请求直接命中。" | 8s |
|
||||
| 4.4 | 高亮 Slot 复用：直接用原 slot 继续写，零拷贝 | "如果原始 slot 空闲，直接原地续写，连 GPU 内存拷贝都不需要。" | 7s |
|
||||
| 4.5 | 首 token 延迟对比：有缓存 vs 无缓存（-50%） | "首 token 延迟降低一半以上。" | 5s |
|
||||
|
||||
**视觉素材**：`prefix_cache.py` 动画、延迟对比
|
||||
|
||||
---
|
||||
|
||||
### Segment 5：Demo 演示（1:50 - 2:15）
|
||||
|
||||
| 镜头 | 画面 | 旁白 | 时长 |
|
||||
|------|------|------|------|
|
||||
| 5.1 | 侧录终端：启动 stream_chat.py，逐行输出对话 | "来实际看看效果。" | 10s |
|
||||
| 5.2 | 多轮对话：中文问答，逐 token 喷出 | 静默 + 打字音效 | 8s |
|
||||
| 5.3 | 切到 HTTP 模式：服务端 + curl 请求，流式返回 | "也提供 OpenAI 兼容的 HTTP API，一行 curl 就能调用。" | 7s |
|
||||
|
||||
**视觉素材**：终端录屏（OBS 录制）
|
||||
|
||||
---
|
||||
|
||||
### Segment 6：收尾 + CTA（2:15 - 2:30）
|
||||
|
||||
| 镜头 | 画面 | 旁白 | 时长 |
|
||||
|------|------|------|------|
|
||||
| 6.1 | 全栈流程回顾（缩略架构图） | "训练用 SEQ → SFT → DPO/GRPO，推理用连续批处理——" | 5s |
|
||||
| 6.2 | GitHub 页面 + Star 引导 | "——全部开源。点个 Star，一起让大模型更普惠。" | 7s |
|
||||
| 6.3 | Logo + URL + "Open Source • Single GPU" | 静默 | 3s |
|
||||
|
||||
**视觉素材**：GitHub 页面录屏、Logo 定版
|
||||
|
||||
---
|
||||
|
||||
## 4. 演示录制指南
|
||||
|
||||
### 4.1 准备工作
|
||||
|
||||
```bash
|
||||
# 1. 安装依赖
|
||||
pip install -e ".[dev]"
|
||||
|
||||
# 2. 下载模型（约 7GB）
|
||||
python scripts/demo/download.py
|
||||
|
||||
# 3. 验证模型加载
|
||||
python scripts/demo/generate_ar.py
|
||||
```
|
||||
|
||||
### 4.2 录制场景 A：交互式对话
|
||||
|
||||
```bash
|
||||
# 终端 1：启动交互式对话
|
||||
python scripts/demo/stream_chat.py
|
||||
|
||||
# 预期交互
|
||||
>> 你好？
|
||||
AstrAI: 你好！有什么我可以帮你的吗？
|
||||
>> 请用中文介绍一下你自己
|
||||
AstrAI: ...（逐 token 输出）
|
||||
>> 编一个关于人工智能的短故事
|
||||
AstrAI: ...（逐 token 输出）
|
||||
```
|
||||
|
||||
**录制重点**：
|
||||
- 逐 token 流式输出效果（用 OBS 录制终端窗口）
|
||||
- 多轮对话的记忆能力（跨轮上下文保持）
|
||||
- 打字音效叠加
|
||||
|
||||
### 4.3 录制场景 B：HTTP 服务 + 并发
|
||||
|
||||
```bash
|
||||
# 终端 1：启动服务器
|
||||
python -m scripts.tools.server --port 8000 --device cuda
|
||||
|
||||
# 终端 2：发送请求（非流式）
|
||||
curl -X POST http://localhost:8000/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"messages":[{"role":"user","content":"Hello!"}],"stream":false}'
|
||||
|
||||
# 终端 3：流式请求
|
||||
curl -X POST http://localhost:8000/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"messages":[{"role":"user","content":"Write a poem"}],"stream":true}'
|
||||
|
||||
# 终端 4：并发压测（用 scripts/demo/generate_batch.py）
|
||||
python scripts/demo/generate_batch.py
|
||||
```
|
||||
|
||||
**录制重点**：
|
||||
- 同时多个 curl 请求展示并发处理
|
||||
- 服务端日志显示批处理合并
|
||||
- `/stats` 端点展示实时统计
|
||||
|
||||
### 4.4 录制规格
|
||||
|
||||
| 参数 | 建议 |
|
||||
|------|------|
|
||||
| 分辨率 | 1920×1080 |
|
||||
| 帧率 | 30fps |
|
||||
| 终端 | Windows Terminal 或 iTerm2，深色主题 |
|
||||
| 字号 | 16-18px，等宽字体（JetBrains Mono / Cascadia Code） |
|
||||
| 录屏工具 | OBS Studio（免费） |
|
||||
| 音频 | 旁白用 USB 麦克风，音效后期叠加 |
|
||||
|
||||
---
|
||||
|
||||
## 5. 动画场景说明
|
||||
|
||||
位于 `promo/` 目录，使用 Manim 引擎。
|
||||
|
||||
### 安装 Manim
|
||||
|
||||
```bash
|
||||
# conda 环境内安装
|
||||
pip install manim
|
||||
|
||||
# 验证
|
||||
python -c "import manim; print(manim.__version__)"
|
||||
```
|
||||
|
||||
### 渲染命令
|
||||
|
||||
```bash
|
||||
# 单独渲染一个场景
|
||||
manim -qh promo/continuous_batching.py ContinuousBatching
|
||||
|
||||
# 全部场景渲染
|
||||
python promo/render_all.py
|
||||
|
||||
# 快速草稿（480p，适合调试）
|
||||
manim -ql promo/continuous_batching.py ContinuousBatching
|
||||
```
|
||||
|
||||
输出文件为 `promo/output/videos/` 下的 `.mp4` 文件，可直接导入剪辑软件。
|
||||
|
||||
### 场景清单
|
||||
|
||||
| 文件 | 导出场景名 | 内容 | 建议时长 |
|
||||
|------|-----------|------|---------|
|
||||
| `transformer.py` | `Transformer` | 模型架构：Embed → GQA → SwiGLU → ×24 → LM Head | ~35s |
|
||||
| `continuous_batching.py` | `ContinuousBatching` | 4 阶段流水线动画 + 吞吐对比 | ~30s |
|
||||
| `prefix_cache.py` | `PrefixCache` | Radix Tree 生长 + 多分支前缀复用 | ~30s |
|
||||
| `architecture.py` | `Architecture` | 全栈架构逐层展开 + 数据流 | ~25s |
|
||||
|
||||
### 自定义动画
|
||||
|
||||
如需修改动画内容：
|
||||
- Manim 语法参考：https://docs.manim.community/
|
||||
- 所有动画元素（颜色、位置、速度）在场景类中通过参数调整
|
||||
- 中文字体渲染需额外配置：
|
||||
|
||||
```python
|
||||
# 在场景类开头添加
|
||||
Text.set_default(font="Microsoft YaHei")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 6. 旁白文案草稿
|
||||
|
||||
### 中文版（完整 2:30）
|
||||
|
||||
```
|
||||
[00:00] 大语言模型很强大——
|
||||
[00:03] 但跑起来需要几十张 GPU，普通人根本碰不到。
|
||||
[00:08] 但如果我告诉你，只要一张显卡就够了呢？
|
||||
[00:13] AstrAI——单卡跑大模型。
|
||||
|
||||
[00:20] 1B 参数，24 层 Transformer，100k 词表的中英 BPE 分词器。
|
||||
[00:28] GQA 分组查询注意力——24 个查询头只对应 4 个 KV 头，KV 缓存直接减少 83%。
|
||||
[00:38] RoPE 旋转位置编码，支持动态长度外推。
|
||||
|
||||
[00:45] 当多个请求同时到达时——
|
||||
[00:48] 传统做法是静态批处理，把请求补齐到相同长度串行处理，GPU 利用率低下。
|
||||
[00:56] AstrAI 采用连续批处理：任务动态进出，GPU 每一刻都在满负荷运转。
|
||||
[01:06] 只有处于相同 KV 缓存位置的任务才一起解码，从根本上避免 RoPE 位置错乱。
|
||||
[01:14] 实测吞吐量提升 3 倍以上。
|
||||
|
||||
[01:20] 如果两个请求有相同的前缀，普通框架会各自从头计算。
|
||||
[01:25] AstrAI 用一颗字典树缓存所有前缀的 KV——第二个请求直接命中。
|
||||
[01:33] 如果原始 slot 空闲，直接原地续写，连 GPU 内存拷贝都不需要。
|
||||
[01:40] 首 token 延迟降低一半以上。
|
||||
|
||||
[01:50] 来实际看看效果。
|
||||
[01:52] （现场演示部分，自由发挥）
|
||||
|
||||
[02:15] 训练到推理，全流程开源，点个 Star，一起让大模型更普惠。
|
||||
[02:25] AstrAI — Open Source, Single GPU.
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 7. 素材清单
|
||||
|
||||
### 视频素材
|
||||
|
||||
| 素材 | 来源 | 状态 |
|
||||
|------|------|------|
|
||||
| 数据中心 / GPU 集群图片 | Pexels / Unsplash 免版权 | 需下载 |
|
||||
| RTX 4090 产品图 | NVIDIA 官网 / 实物拍摄 | 需准备 |
|
||||
| AstrAI Logo | `assets/images/logo.png` | ✅ 已有 |
|
||||
| 终端录屏（对话） | OBS 录制 `scripts/demo/stream_chat.py` | 需录制 |
|
||||
| 终端录屏（HTTP） | OBS 录制 curl + server | 需录制 |
|
||||
| 终端录屏（并发） | OBS 录制 `generate_batch.py` | 需录制 |
|
||||
| GitHub 页面 | 浏览器录屏 | 需录制 |
|
||||
| Transformer 架构动画 | Manim 渲染 `transformer.py` | ✅ 已渲染 |
|
||||
| 架构动画 | Manim 渲染 `architecture.py` | ✅ 已渲染 |
|
||||
| 连续批处理动画 | Manim 渲染 `continuous_batching.py` | ✅ 已渲染 |
|
||||
| 前缀缓存动画 | Manim 渲染 `prefix_cache.py` | ✅ 已渲染 |
|
||||
|
||||
### 音频素材
|
||||
|
||||
| 素材 | 建议 |
|
||||
|------|------|
|
||||
| 旁白 | USB 麦克风录制，男声或女声，中文普通话 |
|
||||
| 背景音乐 | Epidemic Sound / YouTube Audio Library 搜索 "technology ambient" |
|
||||
| 音效 | 打字音效（terminal keystrokes）、转场 swoosh、whoosh |
|
||||
|
||||
### 软件工具
|
||||
|
||||
| 用途 | 推荐工具 | 价格 |
|
||||
|------|---------|------|
|
||||
| 录屏 | OBS Studio | 免费 |
|
||||
| 剪辑 | DaVinci Resolve | 免费 |
|
||||
| 动画渲染 | Manim (`pip install manim`) | 免费 |
|
||||
| 音频处理 | Audacity | 免费 |
|
||||
| 字幕 | DaVinci Resolve 内建 / Aegisub | 免费 |
|
||||
|
||||
---
|
||||
|
||||
## 附录：关键文件索引
|
||||
|
||||
| 文件路径 | 说明 |
|
||||
|---------|------|
|
||||
| `README.md` | 项目主页 README，含快速开始 |
|
||||
| `assets/docs/introduction.md` | 模型架构深度介绍 |
|
||||
| `assets/docs/design.md` | 设计文档 + UML 类图 |
|
||||
| `astrai/inference/scheduler.py` | 连续批处理调度器核心代码 |
|
||||
| `astrai/inference/engine.py` | 推理引擎统一接口 |
|
||||
| `astrai/inference/server.py` | FastAPI 服务器 |
|
||||
| `astrai/model/transformer.py` | Transformer 模型 |
|
||||
| `astrai/model/module.py` | GQA、MLA、MLP 等模块 |
|
||||
| `scripts/demo/stream_chat.py` | 交互式对话演示 |
|
||||
| `scripts/demo/generate_batch.py` | 批量生成演示 |
|
||||
| `scripts/tools/server.py` | HTTP 服务启动脚本 |
|
||||
| `scripts/tools/benchmark.py` | 性能基准测试 |
|
||||
| `scripts/promo/README.md` | 动画渲染说明（已移至 promo/） |
|
||||
| `promo/render_all.py` | 一键渲染所有动画 |
|
||||
| `promo/continuous_batching.py` | 连续批处理 Manim 场景 |
|
||||
| `promo/prefix_cache.py` | 前缀缓存 Manim 场景 |
|
||||
| `promo/architecture.py` | 架构总览 Manim 场景 |
|
||||
| `params/config.json` | 模型配置 |
|
||||
@@ -1,44 +1,270 @@
|
||||
# AstrAI Promo — 动画渲染
|
||||
# AstrAI Promo — 动画渲染 & 视频制作指南
|
||||
|
||||
Manim 驱动的技术动画，用于 AstrAI 宣传视频。
|
||||
目标时长：**2-3 分钟**
|
||||
|
||||
## 依赖
|
||||
---
|
||||
|
||||
```bash
|
||||
pip install manim
|
||||
# 或: conda install -c conda-forge manim
|
||||
## 目录
|
||||
|
||||
1. [项目定位与核心卖点](#1-项目定位与核心卖点)
|
||||
2. [技术架构速览](#2-技术架构速览)
|
||||
3. [渲染指南](#3-渲染指南)
|
||||
4. [场景清单](#4-场景清单)
|
||||
5. [分镜脚本](#5-分镜脚本)
|
||||
6. [旁白文案草稿](#6-旁白文案草稿)
|
||||
7. [素材清单](#7-素材清单)
|
||||
|
||||
---
|
||||
|
||||
## 1. 项目定位与核心卖点
|
||||
|
||||
**一句话定位：**
|
||||
> 一个能在单张消费级 GPU 上训练和推理的 1B 参数中英双语语言模型框架。
|
||||
|
||||
**核心卖点（视频中需突出）：**
|
||||
|
||||
| 卖点 | 说明 | 视觉表达 |
|
||||
|------|------|---------|
|
||||
| **单卡可跑** | 1B 参数，RTX 3090/4090 即可运行 | 巨大服务器集群 vs 单张显卡对比 |
|
||||
| **连续批处理** | 动态合并请求，吞吐量 3x+ | 任务流经 Cleanup→Refill→Prefill→Decode 动画 |
|
||||
| **分页 KV 缓存** | 固定大小页表 + O(1) 分配，按需扩容 | 页表分配与写入动画 |
|
||||
| **OpenAI 兼容 API** | 一行代码切换 | curl 命令对比 |
|
||||
| **流式输出** | 逐 token 返回，低首延迟 | 终端逐字喷出效果 |
|
||||
| **全过程开源** | 训练+推理+权重全部开源 | GitHub 页面展示 |
|
||||
|
||||
---
|
||||
|
||||
## 2. 技术架构速览
|
||||
|
||||
### 整体架构
|
||||
|
||||
```
|
||||
┌──────────────────────────────────────────────────┐
|
||||
│ FastAPI Server (OpenAI-Compatible API) │
|
||||
├──────────────────────────────────────────────────┤
|
||||
│ InferenceEngine (Streaming + Async + Batch) │
|
||||
├──────────────────────────────────────────────────┤
|
||||
│ Continuous Batching Scheduler │
|
||||
│ ┌────────┐ ┌──────┐ ┌────────┐ ┌────────┐ │
|
||||
│ │Cleanup │→ │Refill│→ │Prefill │→ │ Decode │ │
|
||||
│ └────────┘ └──────┘ └────────┘ └────────┘ │
|
||||
├──────────────────────────────────────────────────┤
|
||||
│ Paged KV Cache (Page Table + Page Pool) │
|
||||
├──────────────────────────────────────────────────┤
|
||||
│ Transformer (24层 GQA, RoPE, SwiGLU) │
|
||||
└──────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
## 渲染单个场景
|
||||
### 关键技术指标
|
||||
|
||||
| 指标 | 值 |
|
||||
|------|------|
|
||||
| 参数量 | ~1.0B |
|
||||
| 词表大小 | 100,000（中英 BPE） |
|
||||
| 层数 | 24 |
|
||||
| 注意力头 | 24 Q-heads / 4 KV-heads（GQA） |
|
||||
| 最大长度 | 2048 tokens |
|
||||
| 精度 | bfloat16 |
|
||||
| 最低显存 | ~6GB（推理）/~12GB（训练） |
|
||||
|
||||
---
|
||||
|
||||
## 3. 渲染指南
|
||||
|
||||
依赖 conda 环境 `manim`：
|
||||
|
||||
```bash
|
||||
# 1080p（推荐）
|
||||
manim -qh promo/continuous_batching.py ContinuousBatching
|
||||
|
||||
# 480p 草稿
|
||||
manim -ql promo/continuous_batching.py ContinuousBatching
|
||||
|
||||
# 4K
|
||||
manim -4k promo/continuous_batching.py ContinuousBatching
|
||||
conda run -n manim python -m pip install manim
|
||||
```
|
||||
|
||||
## 渲染全部
|
||||
### 渲染全部（并行）
|
||||
|
||||
```bash
|
||||
python promo/render_all.py
|
||||
conda run -n manim python render_all.py # 1080p
|
||||
conda run -n manim python render_all.py -ql # 480p draft
|
||||
conda run -n manim python render_all.py -4k # 4K
|
||||
```
|
||||
|
||||
输出到 `promo/output/videos/`。
|
||||
### 渲染单个场景
|
||||
|
||||
## 场景清单
|
||||
```bash
|
||||
conda run -n manim python -m manim transformer.py Transformer -qh
|
||||
```
|
||||
|
||||
输出到 `output/videos/`，`.mp4` 可直接导入剪辑软件。
|
||||
|
||||
### 中文字体
|
||||
|
||||
所有 `Text` 对象创建前需设置：
|
||||
|
||||
```python
|
||||
Text.set_default(font="Microsoft YaHei")
|
||||
```
|
||||
|
||||
详细命令参考 `AGENTS.md`。
|
||||
|
||||
---
|
||||
|
||||
## 4. 场景清单
|
||||
|
||||
| 文件 | Scene class | 内容 | 建议时长 |
|
||||
|------|-------------|------|---------|
|
||||
| `transformer.py` | `Transformer` | GQA 注意力机制详解 (Q/K/V→RoPE→Attention→O) + 规格卡 | ~25s |
|
||||
| `continuous_batching.py` | `ContinuousBatching` | 4 阶段流水线 + 吞吐对比 | ~30s |
|
||||
| `prefix_cache.py` | `PrefixCache` | Radix Tree 生长 + 前缀复用 | ~30s |
|
||||
| `architecture.py` | `Architecture` | 全栈架构逐层展开 | ~30s |
|
||||
| `transformer.py` | `Transformer` | GQA 注意力机制 (Q/K/V→RoPE→Attention→O) + 规格卡 | ~35s |
|
||||
| `continuous_batching.py` | `ContinuousBatching` | 4 阶段流水线 + 吞吐对比 (1.0x vs 3.4x) | ~30s |
|
||||
| `paged_cache.py` | `PrefixCache` | 分页 KV 缓存：页表分配、按需扩容、回收 | ~30s |
|
||||
| `architecture.py` | `Architecture` | 全栈架构：HTTP API → Engine → Paged Cache → Transformer | ~25s |
|
||||
| `cta.py` | `CTA` | 结尾卡片："AstrAI — Single GPU, Open Source, 1B params" | ~3s |
|
||||
|
||||
## 导入视频剪辑
|
||||
---
|
||||
|
||||
输出的 `.mp4` 直接导入 DaVinci Resolve / Camtasia / Premiere Pro。
|
||||
## 5. 分镜脚本
|
||||
|
||||
总时长 **2:30**，分为 6 个段落。
|
||||
|
||||
### Segment 1：Hook + 问题陈述（0:00 - 0:20）
|
||||
|
||||
| 镜头 | 画面 | 旁白 | 时长 |
|
||||
|------|------|------|------|
|
||||
| 1.1 | 黑屏，逐字打出"大语言模型很强大" | "大语言模型很强大——" | 3s |
|
||||
| 1.2 | 切到数据中心照片 / 巨大 GPU 集群 | "——但跑起来需要几十张 GPU，普通人根本碰不到。" | 5s |
|
||||
| 1.3 | 画面分屏：左边集群，右边一张 RTX 4090 | "但如果我告诉你，只要一张显卡就够了呢？" | 5s |
|
||||
| 1.4 | Logo 出现：**AstrAI**，下方副标题 "1B 参数单卡推理框架" | "AstrAI——单卡跑大模型。" | 7s |
|
||||
|
||||
**视觉素材**：数据中心图片、RTX 4090 产品图、Logo 动画
|
||||
|
||||
---
|
||||
|
||||
### Segment 2：模型架构速览（0:20 - 0:45）
|
||||
|
||||
| 镜头 | 画面 | 旁白 | 时长 |
|
||||
|------|------|------|------|
|
||||
| 2.1 | Transformer 架构图逐层展开：embed → 24层 decoder → norm → head | "1B 参数，24 层 Transformer，100k 词表的中英 BPE 分词器。" | 8s |
|
||||
| 2.2 | 高亮 GQA：24个 Q head 映射到 4个 KV head | "GQA 分组查询注意力——24 个查询头只对应 4 个 KV 头，KV 缓存直接减少 83%。" | 10s |
|
||||
| 2.3 | RoPE 旋转变换可视化 | "RoPE 旋转位置编码，支持动态长度外推。" | 5s |
|
||||
| 2.4 | fade 到模型 card：vocab=100k, dim=1536, layers=24, heads=24, kv_heads=4 | 静默 | 2s |
|
||||
|
||||
**视觉素材**：`transformer.py` 动画、模型参数 card
|
||||
|
||||
---
|
||||
|
||||
### Segment 3：连续批处理（0:45 - 1:20）
|
||||
|
||||
| 镜头 | 画面 | 旁白 | 时长 |
|
||||
|------|------|------|------|
|
||||
| 3.1 | 3 个请求同时到达服务器 | "当多个请求同时到达时——" | 3s |
|
||||
| 3.2 | 静态批处理对比：最长补齐，3个请求串行 → 总耗时 max_len × 3 | "传统做法是静态批处理，把请求补齐到相同长度，串行处理，GPU 利用率低下。" | 8s |
|
||||
| 3.3 | 连续批处理动画：任务流入 Waiting Queue → Cleanup → Refill → Prefill → Decode | "AstrAI 采用连续批处理：任务动态进出，GPU 每一刻都在满负荷运转。" | 10s |
|
||||
| 3.4 | 放大 Decode 阶段：同一位置的任务合并成一批 | "特别地，只有处于相同 KV 缓存位置的任务才一起解码，从根本上避免了 RoPE 位置错乱的问题。" | 8s |
|
||||
| 3.5 | 吞吐对比柱状图：Static Batch vs Continuous Batching (3x+) | "实测吞吐量提升 3 倍以上。" | 6s |
|
||||
|
||||
**视觉素材**：`continuous_batching.py` 动画、对比图表
|
||||
|
||||
---
|
||||
|
||||
### Segment 4：分页 KV 缓存（1:20 - 1:50）
|
||||
|
||||
| 镜头 | 画面 | 旁白 | 时长 |
|
||||
|------|------|------|------|
|
||||
| 4.1 | 展示 KV 缓存是一个固定大小的张量，被划分为多个相同大小的 page | "KV 缓存不再按请求预分配——而是划分为固定大小的页。" | 5s |
|
||||
| 4.2 | 请求 A 到来，通过页表分配 2 个物理页，写入数据 | "请求到达时，通过页表分配物理页，按需写入。" | 7s |
|
||||
| 4.3 | 请求 B 到来，分配新页，展示页表将逻辑位置映射到不同物理页 | "页表机制让逻辑位置和物理存储解耦——不同请求的页可以分散排列。" | 8s |
|
||||
| 4.4 | Decode 阶段，请求继续生成 token，展示按需分配新页（_maybe_alloc_page） | "生成过程中如果当前页写满，自动追加新页——按需扩容，不浪费显存。" | 7s |
|
||||
| 4.5 | 请求结束时展示页面回收（bitmask 置位） | "请求结束后，页面通过 O(1) 位掩码回收，即刻复用。" | 3s |
|
||||
|
||||
**视觉素材**：`paged_cache.py` 动画、页表分配示意
|
||||
|
||||
---
|
||||
|
||||
### Segment 5：Demo 演示（1:50 - 2:15）
|
||||
|
||||
| 镜头 | 画面 | 旁白 | 时长 |
|
||||
|------|------|------|------|
|
||||
| 5.1 | 侧录终端：启动 stream_chat.py，逐行输出对话 | "来实际看看效果。" | 10s |
|
||||
| 5.2 | 多轮对话：中文问答，逐 token 喷出 | 静默 + 打字音效 | 8s |
|
||||
| 5.3 | 切到 HTTP 模式：服务端 + curl 请求，流式返回 | "也提供 OpenAI 兼容的 HTTP API，一行 curl 就能调用。" | 7s |
|
||||
|
||||
**视觉素材**：终端录屏（OBS 录制）
|
||||
|
||||
---
|
||||
|
||||
### Segment 6：收尾 + CTA（2:15 - 2:30）
|
||||
|
||||
| 镜头 | 画面 | 旁白 | 时长 |
|
||||
|------|------|------|------|
|
||||
| 6.1 | 全栈流程回顾（缩略架构图） | "训练用 SEQ → SFT → DPO/GRPO，推理用连续批处理——" | 5s |
|
||||
| 6.2 | GitHub 页面 + Star 引导 | "——全部开源。点个 Star，一起让大模型更普惠。" | 7s |
|
||||
| 6.3 | Logo + URL + "Open Source • Single GPU" | 静默 | 3s |
|
||||
|
||||
**视觉素材**：GitHub 页面录屏、Logo 定版
|
||||
|
||||
---
|
||||
|
||||
## 6. 旁白文案草稿
|
||||
|
||||
### 中文版（完整 2:30）
|
||||
|
||||
```
|
||||
[00:00] 大语言模型很强大——
|
||||
[00:03] 但跑起来需要几十张 GPU，普通人根本碰不到。
|
||||
[00:08] 但如果我告诉你，只要一张显卡就够了呢？
|
||||
[00:13] AstrAI——单卡跑大模型。
|
||||
|
||||
[00:20] 1B 参数，24 层 Transformer，100k 词表的中英 BPE 分词器。
|
||||
[00:28] GQA 分组查询注意力——24 个查询头只对应 4 个 KV 头，KV 缓存直接减少 83%。
|
||||
[00:38] RoPE 旋转位置编码，支持动态长度外推。
|
||||
|
||||
[00:45] 当多个请求同时到达时——
|
||||
[00:48] 传统做法是静态批处理，把请求补齐到相同长度串行处理，GPU 利用率低下。
|
||||
[00:56] AstrAI 采用连续批处理：任务动态进出，GPU 每一刻都在满负荷运转。
|
||||
[01:06] 只有处于相同 KV 缓存位置的任务才一起解码，从根本上避免 RoPE 位置错乱。
|
||||
[01:14] 实测吞吐量提升 3 倍以上。
|
||||
|
||||
[01:20] 传统 KV 缓存预分配整段显存，浪费严重。
|
||||
[01:25] AstrAI 采用分页 KV 缓存——固定大小的页，通过页表间接寻址，按需分配。
|
||||
[01:33] 生成过程中页写满了自动追加，请求结束后 O(1) 回收。
|
||||
[01:40] 显存利用率大幅提升，支持更多并发请求。
|
||||
|
||||
[01:50] 来实际看看效果。
|
||||
[01:52] （现场演示部分，自由发挥）
|
||||
|
||||
[02:15] 训练到推理，全流程开源，点个 Star，一起让大模型更普惠。
|
||||
[02:25] AstrAI — Open Source, Single GPU.
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 7. 素材清单
|
||||
|
||||
### 视频素材
|
||||
|
||||
| 素材 | 来源 |
|
||||
|------|------|
|
||||
| 数据中心 / GPU 集群图片 | Pexels / Unsplash 免版权 |
|
||||
| RTX 4090 产品图 | NVIDIA 官网 / 实物拍摄 |
|
||||
| AstrAI Logo | 需设计 |
|
||||
| 终端录屏 | OBS 录制 AstrAI demo 脚本 |
|
||||
| GitHub 页面 | 浏览器录屏 |
|
||||
| Transformer 架构动画 | Manim 渲染 `transformer.py` |
|
||||
| 连续批处理动画 | Manim 渲染 `continuous_batching.py` |
|
||||
| 分页缓存动画 | Manim 渲染 `paged_cache.py` |
|
||||
| 全栈架构动画 | Manim 渲染 `architecture.py` |
|
||||
| 结尾卡片 | Manim 渲染 `cta.py` |
|
||||
|
||||
### 音频素材
|
||||
|
||||
| 素材 | 建议 |
|
||||
|------|------|
|
||||
| 旁白 | USB 麦克风录制，中文普通话 |
|
||||
| 背景音乐 | Epidemic Sound / YouTube Audio Library 搜索 "technology ambient" |
|
||||
| 音效 | 打字音效、转场 swoosh |
|
||||
|
||||
### 软件工具
|
||||
|
||||
| 用途 | 推荐工具 | 价格 |
|
||||
|------|---------|------|
|
||||
| 录屏 | OBS Studio | 免费 |
|
||||
| 剪辑 | DaVinci Resolve | 免费 |
|
||||
| 动画渲染 | Manim（conda 环境 `manim`） | 免费 |
|
||||
| 音频处理 | Audacity | 免费 |
|
||||
| 字幕 | DaVinci Resolve 内建 / Aegisub | 免费 |
|
||||
|
||||
+108
-57
@@ -1,76 +1,127 @@
|
||||
"""AstrAI promo: Full architecture overview."""
|
||||
"""AstrAI promo: 4-layer architecture — boxes left, explanations right."""
|
||||
|
||||
from manim import *
|
||||
|
||||
Text.set_default(font="Times New Roman")
|
||||
|
||||
|
||||
class Architecture(Scene):
|
||||
"""Animates the full AstrAI system stack layer by layer."""
|
||||
"""Boxes on left, description text on right for each layer."""
|
||||
|
||||
def construct(self):
|
||||
title = Text("AstrAI Architecture", font_size=48, color=BLUE)
|
||||
title = Text("AstrAI Architecture", font_size=42, color=BLUE)
|
||||
title.to_edge(UP, buff=0.25)
|
||||
self.play(Write(title))
|
||||
self.wait(0.2)
|
||||
self.play(title.animate.to_edge(UP))
|
||||
|
||||
layers_data = [
|
||||
(0.9, GREEN, "API Layer", ["FastAPI Server • OpenAI-Compatible API"]),
|
||||
(0.9, BLUE, "Inference Engine", ["Streaming • Async • Batch Modes"]),
|
||||
(1.6, YELLOW, "Continuous Batching Scheduler",
|
||||
["Cleanup → Refill → Prefill → Decode",
|
||||
"Position-Grouped Decode",
|
||||
"Bitmask O(1) Slot Allocation"]),
|
||||
(1.2, ORANGE, "Prefix Cache + KV Cache",
|
||||
["Radix Tree • Slot Versioning",
|
||||
"GPU copy_() → Zero-Copy Reuse"]),
|
||||
(1.2, PURPLE, "Transformer Model (1B params)",
|
||||
["24-layer GQA • RoPE • SwiGLU",
|
||||
"bfloat16 • 100K vocab"]),
|
||||
W, BH = 5.2, 1.15
|
||||
BX = -3.6
|
||||
TX = 3.6
|
||||
|
||||
def make_box(header, color, bits, src):
|
||||
b = Rectangle(width=W, height=BH, color=color, fill_opacity=0.1, stroke_width=1.5)
|
||||
h = Text(header, font_size=16, color=color)
|
||||
items = [h]
|
||||
for line in bits:
|
||||
items.append(Text(line, font_size=10, color=WHITE))
|
||||
items.append(Text(src, font_size=9, color=GRAY))
|
||||
c = VGroup(*items).arrange(DOWN, buff=0.04)
|
||||
c.move_to(b.get_center())
|
||||
return VGroup(b, c)
|
||||
|
||||
L1 = make_box("HTTP API Server", GREEN,
|
||||
["FastAPI · OpenAI-Compatible",
|
||||
"/v1/chat/completions · SSE streaming"],
|
||||
"astrai/inference/server.py")
|
||||
|
||||
L2 = make_box("Inference Engine", BLUE,
|
||||
["generate() · batch mode · streaming",
|
||||
"4-phase daemon: Cleanup → Refill → Prefill → Decode",
|
||||
"Position-grouped decode · Bitmask O(1) slots"],
|
||||
"astrai/inference/engine.py · scheduler.py")
|
||||
|
||||
L3 = make_box("Paged KV Cache", ORANGE,
|
||||
["Page-table-indirected read/write · Page pool",
|
||||
"O(1) bitmask alloc/free · On-demand page growth"],
|
||||
"astrai/inference/cache.py · scheduler.py")
|
||||
|
||||
L4 = make_box("Transformer Model", PURPLE,
|
||||
["24× DecoderBlock · GQA 6:1 · RoPE",
|
||||
"SwiGLU MLP · Dim 1536 · bfloat16"],
|
||||
"astrai/model/transformer.py")
|
||||
|
||||
layers = VGroup(L1, L2, L3, L4)
|
||||
layers.arrange(DOWN, buff=0.08)
|
||||
layers.move_to([BX, 0, 0])
|
||||
layers.next_to(title, DOWN, buff=0.25)
|
||||
|
||||
# Description panels (right side)
|
||||
descs_text = [
|
||||
["HTTP API Server",
|
||||
"Receives chat requests via",
|
||||
"OpenAI-compatible endpoints.",
|
||||
"Streams generated tokens back",
|
||||
"through Server-Sent Events."],
|
||||
["Inference Engine",
|
||||
"Orchestrates the full generation",
|
||||
"pipeline with a background daemon.",
|
||||
"4-phase loop: Cleanup tasks,",
|
||||
"Refill batch, Prefill prompts,",
|
||||
"Decode tokens one by one."],
|
||||
["Paged KV Cache",
|
||||
"Divides KV cache into fixed-size pages",
|
||||
"with page-table-indirected access.",
|
||||
"Per-task page tables map logical pages",
|
||||
"to physical pages — O(1) alloc/free."],
|
||||
["Transformer Model",
|
||||
"Decoder-only Transformer with",
|
||||
"Grouped-Query Attention (GQA 6:1).",
|
||||
"RoPE rotary encoding, SwiGLU",
|
||||
"activation, 100K vocabulary."],
|
||||
]
|
||||
|
||||
layers = VGroup()
|
||||
for height, color, label, subs in layers_data:
|
||||
box = Rectangle(width=7.5, height=height, color=color, fill_opacity=0.1)
|
||||
lbl = Text(label, font_size=18, color=color)
|
||||
items = [lbl] + [Text(s, font_size=11, color=WHITE) for s in subs]
|
||||
content = VGroup(*items)
|
||||
content.arrange(DOWN, buff=0.22)
|
||||
content.move_to(box.get_center())
|
||||
layers.add(VGroup(box, content))
|
||||
def make_desc(lines, color):
|
||||
els = [Text(lines[0], font_size=20, color=color)]
|
||||
for ln in lines[1:]:
|
||||
els.append(Text(ln, font_size=14, color=WHITE))
|
||||
grp = VGroup(*els).arrange(DOWN, buff=0.1, aligned_edge=LEFT)
|
||||
return grp
|
||||
|
||||
layers.arrange(DOWN, buff=0.18)
|
||||
layers.next_to(title, DOWN, buff=0.3)
|
||||
COLORS = [GREEN, BLUE, ORANGE, PURPLE]
|
||||
descs = [make_desc(lns, c) for lns, c in zip(descs_text, COLORS)]
|
||||
|
||||
for i in range(len(layers)):
|
||||
self.play(Create(layers[i]), run_time=0.35)
|
||||
if i > 0:
|
||||
# Use box-to-box for arrow endpoints (not content)
|
||||
prev_box = layers[i - 1][0]
|
||||
curr_box = layers[i][0]
|
||||
arrows = VGroup()
|
||||
for i, (layer, desc) in enumerate(zip(layers, descs)):
|
||||
b = layer[0]
|
||||
self.play(Create(layer), run_time=0.35)
|
||||
desc.next_to(b, RIGHT, buff=1.0)
|
||||
desc.align_to(b, UP)
|
||||
self.play(Write(desc), run_time=0.3)
|
||||
self.wait(2.0 if i == 0 else 1.8)
|
||||
if i < len(layers) - 1:
|
||||
self.play(FadeOut(desc))
|
||||
nxt = layers[i + 1][0]
|
||||
arrow = Arrow(
|
||||
prev_box.get_bottom(),
|
||||
curr_box.get_top(),
|
||||
color=GRAY,
|
||||
buff=0.1,
|
||||
max_tip_length_to_length_ratio=0.15,
|
||||
b.get_bottom(), nxt.get_top(),
|
||||
color=GRAY, buff=0.04,
|
||||
max_tip_length_to_length_ratio=0.18,
|
||||
)
|
||||
self.play(Create(arrow), run_time=0.15)
|
||||
|
||||
self.play(Create(arrow), run_time=0.12)
|
||||
arrows.add(arrow)
|
||||
else:
|
||||
self.wait(0.5)
|
||||
self.play(FadeOut(desc))
|
||||
|
||||
hl = SurroundingRectangle(layers[3], color=GREEN, buff=0.12)
|
||||
hl_note = Text("Zero-Copy Prefix Reuse", font_size=18, color=GREEN)
|
||||
hl_note.next_to(hl, LEFT, buff=0.4)
|
||||
self.play(Create(hl), Write(hl_note))
|
||||
self.wait(1.5)
|
||||
self.play(FadeOut(hl), FadeOut(hl_note))
|
||||
# Show all boxes + arrows together briefly
|
||||
self.wait(0.3)
|
||||
|
||||
self.play(FadeOut(layers))
|
||||
# Highlight innovation layers
|
||||
hl2 = SurroundingRectangle(L2, color=BLUE, buff=0.1, stroke_width=2)
|
||||
hl3 = SurroundingRectangle(L3, color=ORANGE, buff=0.1, stroke_width=2)
|
||||
hl_note = Text("Key Innovations", font_size=20, color=GOLD)
|
||||
hl_note.next_to(VGroup(hl2, hl3), RIGHT, buff=1.5)
|
||||
hl_note.align_to(hl2, UP)
|
||||
self.play(Create(hl2), Create(hl3), Write(hl_note))
|
||||
self.wait(2.0)
|
||||
self.play(FadeOut(hl2), FadeOut(hl3), FadeOut(hl_note))
|
||||
|
||||
cta = VGroup(
|
||||
Text("AstrAI", font_size=52, color=BLUE),
|
||||
Text("Single GPU • Open Source • 1B params", font_size=24, color=GRAY),
|
||||
Text("github.com/ViperEkura/AstrAI", font_size=20, color=YELLOW),
|
||||
).arrange(DOWN, buff=0.35)
|
||||
self.play(Write(cta))
|
||||
self.wait(2)
|
||||
self.play(FadeOut(cta), FadeOut(title))
|
||||
self.play(FadeOut(VGroup(*layers)), FadeOut(arrows), FadeOut(title))
|
||||
|
||||
+322
-196
@@ -6,8 +6,9 @@ with coloured batch tokens flowing through states, entering & leaving continuous
|
||||
|
||||
from manim import *
|
||||
|
||||
# ── palette ──
|
||||
PHASE_COLORS = {
|
||||
Text.set_default(font="Times New Roman")
|
||||
|
||||
PAL = {
|
||||
"Cleanup": GRAY,
|
||||
"Refill": ORANGE,
|
||||
"Prefill": BLUE,
|
||||
@@ -28,217 +29,169 @@ class ContinuousBatching(Scene):
|
||||
bar = Line(LEFT * 7, RIGHT * 7, color=GRAY).next_to(title, DOWN, buff=0.15)
|
||||
self.play(Create(bar))
|
||||
|
||||
# ═══════════════════════════════════════════════════
|
||||
# 1. Build state-machine layout (vertical, 4 states)
|
||||
# ═══════════════════════════════════════════════════
|
||||
state_names = ["Cleanup", "Refill", "Prefill", "Decode"]
|
||||
# ── layout config ──
|
||||
LANE_W, LANE_H = 2.6, 0.95
|
||||
X_P, X_Pr, X_D, X_F = -4.95, -1.65, 1.65, 4.95
|
||||
YL = 0.3
|
||||
P_CLR, F_CLR = GRAY, RED
|
||||
|
||||
states = VGroup()
|
||||
trans_arrows = VGroup()
|
||||
for i, name in enumerate(state_names):
|
||||
def lane(x, label, clr, sub):
|
||||
box = RoundedRectangle(
|
||||
width=3.6, height=0.8, corner_radius=0.15,
|
||||
color=PHASE_COLORS[name], fill_opacity=0.12, stroke_width=2.5,
|
||||
width=LANE_W, height=LANE_H, corner_radius=0.12,
|
||||
color=clr, fill_opacity=0.10, stroke_width=2.2,
|
||||
)
|
||||
lbl = Text(name, font_size=20, color=PHASE_COLORS[name])
|
||||
states.add(VGroup(box, lbl))
|
||||
t = Text(label, font_size=20, color=clr)
|
||||
s = Text(sub, font_size=10, color=LIGHT_GRAY)
|
||||
inner = VGroup(t, s).arrange(DOWN, buff=0.04).move_to(box)
|
||||
return VGroup(box, inner).move_to([x, YL, 0])
|
||||
|
||||
states.arrange(DOWN, buff=0.3)
|
||||
states.shift(LEFT * 3.8 + DOWN * 0.5)
|
||||
# ── FSM label (unchanged) ──
|
||||
fsm_states = VGroup()
|
||||
for label, clr in [("Refill", ORANGE), ("→", LIGHT_GRAY),
|
||||
("Prefill", BLUE), ("→", LIGHT_GRAY),
|
||||
("Decode", YELLOW), ("→", LIGHT_GRAY),
|
||||
("Cleanup", GRAY)]:
|
||||
t = Text(label, font_size=13, color=clr)
|
||||
fsm_states.add(t)
|
||||
fsm_states.arrange(RIGHT, buff=0.06)
|
||||
fsm_states.next_to(bar, DOWN, buff=0.3)
|
||||
|
||||
for i in range(1, 4):
|
||||
a = Arrow(
|
||||
states[i - 1].get_bottom(), states[i].get_top(),
|
||||
# ── 4-state pipeline lanes ──
|
||||
pend_lane = lane(X_P, "PENDING", P_CLR, "waiting queue")
|
||||
pref_lane = lane(X_Pr, "PREFILL", BLUE, "first token")
|
||||
dec_lane = lane(X_D, "DECODE", YELLOW, "per-token gen")
|
||||
fin_lane = lane(X_F, "FINISHED", F_CLR, "sequence done")
|
||||
lane_group = VGroup(pend_lane, pref_lane, dec_lane, fin_lane)
|
||||
|
||||
# ── arrows ──
|
||||
ea = Arrow(pend_lane.get_left() + LEFT * 0.9, pend_lane.get_left(),
|
||||
color=GREEN, stroke_width=2.5,
|
||||
max_tip_length_to_length_ratio=0.15)
|
||||
el = Text("New Req", font_size=11, color=GREEN)
|
||||
el.next_to(ea, UP, buff=0.04)
|
||||
|
||||
ra = Arrow(pend_lane.get_right(), pref_lane.get_left(),
|
||||
color=ORANGE, buff=0.06,
|
||||
max_tip_length_to_length_ratio=0.15)
|
||||
rl = Text("Refill", font_size=10, color=ORANGE)
|
||||
rl.next_to(ra, UP, buff=0.04)
|
||||
|
||||
ta = Arrow(pref_lane.get_right(), dec_lane.get_left(),
|
||||
color=LIGHT_GRAY, buff=0.06,
|
||||
max_tip_length_to_length_ratio=0.22,
|
||||
)
|
||||
trans_arrows.add(a)
|
||||
max_tip_length_to_length_ratio=0.15)
|
||||
|
||||
for i in range(4):
|
||||
self.play(Create(states[i]))
|
||||
if i > 0:
|
||||
self.play(Create(trans_arrows[i - 1]))
|
||||
ca = Arrow(dec_lane.get_right(), fin_lane.get_left(),
|
||||
color=GRAY, buff=0.06,
|
||||
max_tip_length_to_length_ratio=0.15)
|
||||
cl = Text("Cleanup", font_size=10, color=GRAY)
|
||||
cl.next_to(ca, UP, buff=0.04)
|
||||
|
||||
# loop arrow — Decode returns to Cleanup (multiturn decoding)
|
||||
loop = CurvedArrow(
|
||||
states[-1].get_right() + RIGHT * 0.2,
|
||||
states[0].get_right() + RIGHT * 0.2,
|
||||
color=LIGHT_GRAY, angle=PI / 2,
|
||||
)
|
||||
loop_lbl = Text("per token", font_size=11, color=GRAY).next_to(loop, RIGHT, buff=0.08)
|
||||
self.play(Create(loop), Write(loop_lbl))
|
||||
self.wait(0.4)
|
||||
xa = Arrow(fin_lane.get_right(), fin_lane.get_right() + RIGHT * 0.9,
|
||||
color=RED, stroke_width=2.5,
|
||||
max_tip_length_to_length_ratio=0.15)
|
||||
xl = Text("Exit", font_size=11, color=RED)
|
||||
xl.next_to(xa, UP, buff=0.04)
|
||||
|
||||
# ═══════════════════════════════════════════════════
|
||||
# 2. Boot tokens — initial batches placed at mid-cycle
|
||||
# ═══════════════════════════════════════════════════
|
||||
def make_token(name: str, col: str) -> VGroup:
|
||||
card = RoundedRectangle(width=0.65, height=0.38, corner_radius=0.08,
|
||||
color=col, fill_opacity=0.35, stroke_width=1.8)
|
||||
txt = Text(name, font_size=13, color=col)
|
||||
return VGroup(card, txt)
|
||||
|
||||
tokens = {
|
||||
"A": make_token("A", BATCH_COLORS[0]),
|
||||
"B": make_token("B", BATCH_COLORS[1]),
|
||||
"C": make_token("C", BATCH_COLORS[2]),
|
||||
}
|
||||
# all three at consecutive stages, Prefill is the entry point
|
||||
tokens["A"].move_to(states[2]).shift(RIGHT * 1.5) # Prefill
|
||||
tokens["B"].move_to(states[3]).shift(RIGHT * 1.5) # Decode
|
||||
tokens["C"].move_to(states[0]).shift(RIGHT * 1.5) # Cleanup
|
||||
|
||||
for t in tokens.values():
|
||||
self.play(FadeIn(t, scale=0.7), run_time=0.25)
|
||||
self.wait(0.2)
|
||||
|
||||
note = Text("Every request starts at Prefill", font_size=16, color=WHITE) \
|
||||
.next_to(states, DOWN, buff=0.55)
|
||||
self.play(Write(note))
|
||||
self.wait(1.0)
|
||||
self.play(FadeOut(note))
|
||||
|
||||
# ═══════════════════════════════════════════════════
|
||||
# 3. Tick 1 — advance, C exits, new D enters at Prefill
|
||||
# ═══════════════════════════════════════════════════
|
||||
slots = [
|
||||
states[0].get_center() + RIGHT * 1.5, # Cleanup
|
||||
states[1].get_center() + RIGHT * 1.5, # Refill
|
||||
states[2].get_center() + RIGHT * 1.5, # Prefill
|
||||
states[3].get_center() + RIGHT * 1.5, # Decode
|
||||
]
|
||||
|
||||
self.play(
|
||||
tokens["A"].animate.move_to(slots[3]), # Prefill → Decode
|
||||
tokens["B"].animate.move_to(slots[0]), # Decode → Cleanup
|
||||
tokens["C"].animate.move_to(slots[1]), # Cleanup → Refill
|
||||
)
|
||||
self.play(Write(fsm_states))
|
||||
self.play(Create(pend_lane), Create(pref_lane), Create(dec_lane), Create(fin_lane))
|
||||
self.wait(0.3)
|
||||
self.play(Create(ea), Write(el),
|
||||
Create(ra), Write(rl),
|
||||
Create(ta),
|
||||
Create(ca), Write(cl),
|
||||
Create(xa), Write(xl))
|
||||
self.wait(0.5)
|
||||
|
||||
# C (now at Refill) exits after completing the loop
|
||||
# new D enters at Prefill
|
||||
self.play(FadeOut(tokens["C"], scale=0.6))
|
||||
tokens["D"] = make_token("D", BATCH_COLORS[3])
|
||||
tokens["D"].move_to(states[2]).shift(RIGHT * 1.5) # Prefill ← entry
|
||||
self.play(FadeIn(tokens["D"], scale=0.7))
|
||||
self.wait(0.25)
|
||||
# ── Tokens ──
|
||||
TOK_W, TOK_H = 0.58, 0.38
|
||||
|
||||
# ═══════════════════════════════════════════════════
|
||||
# 4. Tick 2 — advance, B exits, new E enters at Prefill
|
||||
# ═══════════════════════════════════════════════════
|
||||
self.play(
|
||||
tokens["D"].animate.move_to(slots[3]), # Prefill → Decode
|
||||
tokens["A"].animate.move_to(slots[0]), # Decode → Cleanup
|
||||
tokens["B"].animate.move_to(slots[1]), # Cleanup → Refill
|
||||
def mk_tok(name, col, state, n_tok):
|
||||
card = RoundedRectangle(
|
||||
width=TOK_W, height=TOK_H, corner_radius=0.06,
|
||||
color=col, fill_opacity=0.38, stroke_width=1.6,
|
||||
)
|
||||
self.wait(0.3)
|
||||
t = Text(name, font_size=13, color=col).move_to(card)
|
||||
info = Text(f"{state} {n_tok}t", font_size=7, color=col)
|
||||
return VGroup(VGroup(card, t), info).arrange(DOWN, buff=0.03)
|
||||
|
||||
self.play(FadeOut(tokens["B"], scale=0.6))
|
||||
tokens["E"] = make_token("E", BATCH_COLORS[4])
|
||||
tokens["E"].move_to(states[2]).shift(RIGHT * 1.5) # Prefill ← entry
|
||||
self.play(FadeIn(tokens["E"], scale=0.7))
|
||||
self.wait(0.25)
|
||||
def slots(x, n):
|
||||
sp = LANE_W * 0.72 / max(n, 1)
|
||||
sx = x - (n - 1) * sp / 2
|
||||
return [np.array([sx + i * sp, -1.5, 0]) for i in range(n)]
|
||||
|
||||
P_SLOTS = slots(X_P, 2) # G, F
|
||||
Pr_SLOTS = slots(X_Pr, 1) # E
|
||||
D_SLOTS = slots(X_D, 3) # D, A, B
|
||||
F_SLOTS = slots(X_F, 1) # C
|
||||
|
||||
tok = {}
|
||||
def add(name, col, lane_slots, idx, state, n):
|
||||
t = mk_tok(name, col, state, n).move_to(lane_slots[idx])
|
||||
tok[name] = t
|
||||
|
||||
add("G", BATCH_COLORS[6], P_SLOTS, 0, "PENDING", 0)
|
||||
add("F", BATCH_COLORS[5], P_SLOTS, 1, "PENDING", 0)
|
||||
add("E", BATCH_COLORS[4], Pr_SLOTS, 0, "PREFILL", 128)
|
||||
add("D", BATCH_COLORS[3], D_SLOTS, 0, "DECODE", 5)
|
||||
add("A", BATCH_COLORS[0], D_SLOTS, 1, "DECODE", 9)
|
||||
add("B", BATCH_COLORS[1], D_SLOTS, 2, "DECODE", 13)
|
||||
add("C", BATCH_COLORS[2], F_SLOTS, 0, "FINISHED", 16)
|
||||
|
||||
for t in tok.values():
|
||||
self.play(FadeIn(t, scale=0.7), run_time=0.18)
|
||||
self.wait(2.0)
|
||||
|
||||
self.play(*[FadeOut(t) for t in tok.values()])
|
||||
|
||||
# ═══════════════════════════════════════════════════
|
||||
# 5. Tick 3 — advance, A exits, new F enters at Prefill
|
||||
# 7. Position-Grouped Decode highlight
|
||||
# ═══════════════════════════════════════════════════
|
||||
self.play(
|
||||
tokens["E"].animate.move_to(slots[3]), # Prefill → Decode
|
||||
tokens["D"].animate.move_to(slots[0]), # Decode → Cleanup
|
||||
tokens["A"].animate.move_to(slots[1]), # Cleanup → Refill
|
||||
)
|
||||
self.wait(0.25)
|
||||
|
||||
self.play(FadeOut(tokens["A"], scale=0.6))
|
||||
tokens["F"] = make_token("F", BATCH_COLORS[5])
|
||||
tokens["F"].move_to(states[2]).shift(RIGHT * 1.5) # Prefill ← entry
|
||||
self.play(FadeIn(tokens["F"], scale=0.7))
|
||||
self.wait(0.25)
|
||||
|
||||
# ═══════════════════════════════════════════════════
|
||||
# 6. Tick 4 — advance, F exits, new G enters at Prefill
|
||||
# ═══════════════════════════════════════════════════
|
||||
self.play(
|
||||
tokens["F"].animate.move_to(slots[3]), # Prefill → Decode
|
||||
tokens["E"].animate.move_to(slots[0]), # Decode → Cleanup
|
||||
tokens["D"].animate.move_to(slots[1]), # Cleanup → Refill
|
||||
)
|
||||
self.wait(0.25)
|
||||
|
||||
self.play(FadeOut(tokens["D"], scale=0.6))
|
||||
tokens["G"] = make_token("G", BATCH_COLORS[6])
|
||||
tokens["G"].move_to(states[2]).shift(RIGHT * 1.5) # Prefill ← entry
|
||||
self.play(FadeIn(tokens["G"], scale=0.7))
|
||||
self.wait(0.35)
|
||||
|
||||
# drop note: constant throughput, all enter at Prefill
|
||||
flow_note = Text("All requests enter at Prefill — pipeline never drains",
|
||||
font_size=15, color=GREEN).next_to(states, DOWN, buff=0.55)
|
||||
self.play(Write(flow_note))
|
||||
self.wait(1.5)
|
||||
self.play(FadeOut(flow_note))
|
||||
|
||||
# clear tokens
|
||||
self.play(*[FadeOut(t) for t in tokens.values()])
|
||||
|
||||
# ═══════════════════════════════════════════════════
|
||||
# 6. Position-Grouped Decode highlight
|
||||
# ═══════════════════════════════════════════════════
|
||||
# show multiple tokens grouped at Decode
|
||||
d_pos = states[3].get_center()
|
||||
d_tokens = [
|
||||
make_token("T" + str(i), BATCH_COLORS[i]) for i in range(4)
|
||||
]
|
||||
positions = [
|
||||
d_pos + RIGHT * 1.2 + UP * 0.45,
|
||||
d_pos + RIGHT * 1.2,
|
||||
d_pos + RIGHT * 2.5 + UP * 0.45,
|
||||
d_pos + RIGHT * 2.5,
|
||||
]
|
||||
for i in range(4):
|
||||
d_tokens[i].move_to(positions[i])
|
||||
self.play(FadeIn(d_tokens[i], scale=0.6), run_time=0.2)
|
||||
|
||||
ring = SurroundingRectangle(states[3], color=YELLOW, buff=0.12, stroke_width=3)
|
||||
ring = SurroundingRectangle(dec_lane, color=YELLOW, buff=0.12, stroke_width=3)
|
||||
ring_txt = Text(
|
||||
"Position-Grouped Batching\nSame decode position → single matmul",
|
||||
font_size=14, color=YELLOW, line_spacing=0.6,
|
||||
).next_to(states[3], DOWN, buff=0.5)
|
||||
).next_to(dec_lane, DOWN, buff=0.5)
|
||||
self.play(Create(ring), Write(ring_txt))
|
||||
self.wait(2.0)
|
||||
self.play(FadeOut(ring), FadeOut(ring_txt),
|
||||
*[FadeOut(t) for t in d_tokens])
|
||||
self.play(FadeOut(ring), FadeOut(ring_txt))
|
||||
|
||||
# ═══════════════════════════════════════════════════
|
||||
# 7. O(1) Bitmask Slot Allocation
|
||||
# 8. O(1) Bitmask Slot Allocation
|
||||
# ═══════════════════════════════════════════════════
|
||||
bitmask_title = Text("O(1) Slot Allocation via Bitmask",
|
||||
font_size=22, color=ORANGE).next_to(states, DOWN, buff=0.75)
|
||||
bitmask_desc = Text("free_slots = ~occupied_mask (one-clock op)",
|
||||
font_size=15, color=GRAY).next_to(bitmask_title, DOWN, buff=0.15)
|
||||
bitmask_title = Text(
|
||||
"O(1) Slot Allocation via Bitmask",
|
||||
font_size=22, color=ORANGE,
|
||||
).next_to(lane_group, DOWN, buff=0.75)
|
||||
bitmask_desc = Text(
|
||||
"free_slots = ~occupied_mask (one-clock op)",
|
||||
font_size=15, color=GRAY,
|
||||
).next_to(bitmask_title, DOWN, buff=0.15)
|
||||
self.play(Write(bitmask_title), Write(bitmask_desc))
|
||||
self.wait(1.5)
|
||||
|
||||
# animate bitmask bits flipping
|
||||
bits_group = VGroup()
|
||||
bit_size = 0.18
|
||||
for i in range(16):
|
||||
square = Square(side_length=bit_size * 2, color=GRAY,
|
||||
fill_opacity=0.0, stroke_width=1.2)
|
||||
square = Square(
|
||||
side_length=bit_size * 2, color=GRAY,
|
||||
fill_opacity=0.0, stroke_width=1.2,
|
||||
)
|
||||
if i in (2, 5, 9, 13):
|
||||
square.set_fill(GRAY, opacity=0.5)
|
||||
bits_group.add(square)
|
||||
bits_group.arrange(RIGHT, buff=0.06)
|
||||
bits_group.next_to(bitmask_desc, DOWN, buff=0.3)
|
||||
|
||||
occupied_lbl = Text("occupied_mask", font_size=11, color=RED).next_to(bits_group, LEFT, buff=0.4)
|
||||
occupied_lbl = Text("occupied_mask", font_size=11, color=RED) \
|
||||
.next_to(bits_group, LEFT, buff=0.4)
|
||||
self.play(Create(bits_group), Write(occupied_lbl))
|
||||
|
||||
# flip to ~occupied
|
||||
flipped = VGroup()
|
||||
for i, sq in enumerate(bits_group):
|
||||
copy_sq = Square(side_length=bit_size * 2, color=GRAY,
|
||||
fill_opacity=0.0, stroke_width=1.2).move_to(sq)
|
||||
copy_sq = Square(
|
||||
side_length=bit_size * 2, color=GRAY,
|
||||
fill_opacity=0.0, stroke_width=1.2,
|
||||
).move_to(sq)
|
||||
if i not in (2, 5, 9, 13):
|
||||
copy_sq.set_fill(GRAY, opacity=0.5)
|
||||
flipped.add(copy_sq)
|
||||
@@ -252,55 +205,230 @@ class ContinuousBatching(Scene):
|
||||
FadeOut(bitmask_title), FadeOut(bitmask_desc))
|
||||
|
||||
# ═══════════════════════════════════════════════════
|
||||
# 8. Throughput comparison with animated bars
|
||||
# 9. Gantt timeline comparison — Static vs Continuous
|
||||
# ═══════════════════════════════════════════════════
|
||||
self.play(
|
||||
*[FadeOut(m) for m in self.mobjects if m is not title and m is not bar],
|
||||
FadeOut(loop), FadeOut(loop_lbl),
|
||||
)
|
||||
for s in states:
|
||||
self.play(FadeOut(s), run_time=0.15)
|
||||
for a in trans_arrows:
|
||||
self.play(FadeOut(a), run_time=0.15)
|
||||
self.wait(0.2)
|
||||
|
||||
CELL = 0.44
|
||||
BH = 0.32
|
||||
BGAP = 0.10
|
||||
ROW = BH + BGAP
|
||||
TICKS = 12
|
||||
PANEL_W = TICKS * CELL
|
||||
L_OX = -5.8
|
||||
R_OX = 1.0
|
||||
GY = 2.0
|
||||
|
||||
def gbox(ox, y, start, span, color, fill=0.75):
|
||||
x = ox + start * CELL
|
||||
w = span * CELL
|
||||
return Rectangle(
|
||||
width=w, height=BH, color=color,
|
||||
fill_opacity=fill, stroke_width=0,
|
||||
).move_to([x + w / 2, y, 0])
|
||||
|
||||
def batch_box(ox, y_gpu, y_last_req, start, span, color, label_txt):
|
||||
w = span * CELL
|
||||
top = y_gpu + BH / 2 + 0.06
|
||||
bot = y_last_req - BH / 2 - 0.06
|
||||
h = top - bot
|
||||
cx = ox + (start + span / 2) * CELL
|
||||
cy = (top + bot) / 2
|
||||
rect = Rectangle(
|
||||
width=w, height=h, color=color,
|
||||
stroke_width=1.8, fill_opacity=0.04,
|
||||
)
|
||||
rect.move_to([cx, cy, 0])
|
||||
lbl = Text(label_txt, font_size=12, color=color).next_to(rect, UP, buff=0.06)
|
||||
return rect, lbl
|
||||
|
||||
def taxis(ox, ty):
|
||||
line = Line(
|
||||
[ox, ty, 0], [ox + PANEL_W, ty, 0],
|
||||
color=GRAY, stroke_width=1.2,
|
||||
)
|
||||
ticks_vg = VGroup()
|
||||
for t in range(TICKS + 1):
|
||||
ti = Line(DOWN * 0.06, UP * 0.06, color=GRAY, stroke_width=0.8)
|
||||
ti.move_to([ox + t * CELL, ty, 0])
|
||||
ticks_vg.add(ti)
|
||||
nums_vg = VGroup()
|
||||
for t in range(0, TICKS + 1, 3):
|
||||
n = Text(str(t), font_size=11, color=GRAY).next_to(
|
||||
[ox + t * CELL, ty, 0], DOWN, buff=0.10,
|
||||
)
|
||||
nums_vg.add(n)
|
||||
return VGroup(line, ticks_vg, nums_vg)
|
||||
|
||||
# ── Left: Static Batching ──
|
||||
s_title = Text("Static Batching", font_size=26, color=RED)
|
||||
s_title.move_to([L_OX + PANEL_W / 2, GY + 0.65, 0])
|
||||
s_note = Text(
|
||||
"requests wait → batch together → all run same length · GPU idle gaps",
|
||||
font_size=13, color=RED,
|
||||
).move_to([L_OX + PANEL_W / 2, -1.6, 0])
|
||||
self.play(Write(s_title))
|
||||
self.wait(0.25)
|
||||
|
||||
st_axis = taxis(L_OX, GY)
|
||||
self.play(Create(st_axis))
|
||||
|
||||
gpu_l = Text("GPU", font_size=14, color=WHITE)
|
||||
gpu_l.move_to([L_OX - 0.55, GY - ROW, 0])
|
||||
self.play(Write(gpu_l))
|
||||
|
||||
s_y_gpu = GY - ROW
|
||||
s_gpu_idle1 = gbox(L_OX, s_y_gpu, 0, 2, RED, 0.45)
|
||||
s_gpu_batch1 = gbox(L_OX, s_y_gpu, 2, 4, GREEN)
|
||||
s_gpu_idle2 = gbox(L_OX, s_y_gpu, 6, 2, RED, 0.45)
|
||||
s_gpu_batch2 = gbox(L_OX, s_y_gpu, 8, 4, GREEN)
|
||||
s_gpu_bars = [s_gpu_idle1, s_gpu_batch1, s_gpu_idle2, s_gpu_batch2]
|
||||
for seg in s_gpu_bars:
|
||||
self.play(GrowFromEdge(seg, LEFT), run_time=0.09)
|
||||
|
||||
s_idle1 = Text("IDLE", font_size=10, color=RED) \
|
||||
.move_to([L_OX + 1 * CELL, s_y_gpu, 0])
|
||||
s_idle2 = Text("IDLE", font_size=10, color=RED) \
|
||||
.move_to([L_OX + 7 * CELL, s_y_gpu, 0])
|
||||
self.play(Write(s_idle1), Write(s_idle2))
|
||||
|
||||
s_req_defs = [
|
||||
("A", ORANGE, 0, 2, 2, 6),
|
||||
("B", BLUE, 1, 2, 2, 6),
|
||||
("C", PINK, 2, 2, 2, 6),
|
||||
("D", ORANGE, 4, 8, 8, 12),
|
||||
("E", BLUE, 6, 8, 8, 12),
|
||||
("F", PINK, 8, 8, 8, 12),
|
||||
]
|
||||
s_bars = []
|
||||
for i, (name, col, ws, we, rs, re) in enumerate(s_req_defs):
|
||||
y = s_y_gpu - (i + 1) * ROW
|
||||
lbl = Text(f"Req {name}", font_size=12, color=col)
|
||||
lbl.move_to([L_OX - 0.55, y, 0])
|
||||
items = [lbl]
|
||||
anims = [FadeIn(lbl)]
|
||||
if we - ws > 0.02:
|
||||
wbar = gbox(L_OX, y, ws, we - ws, GRAY, 0.28)
|
||||
items.append(wbar)
|
||||
anims.append(GrowFromEdge(wbar, LEFT))
|
||||
rbar = gbox(L_OX, y, rs, re - rs, col, 0.60)
|
||||
items.append(rbar)
|
||||
anims.append(GrowFromEdge(rbar, LEFT))
|
||||
s_bars.extend(items)
|
||||
self.play(*anims, run_time=0.09)
|
||||
|
||||
s_y_last3 = s_y_gpu - 3 * ROW
|
||||
s_y_last6 = s_y_gpu - 6 * ROW
|
||||
b1_rect, b1_lbl = batch_box(L_OX, s_y_gpu, s_y_last3, 2, 4, RED, "Batch 1")
|
||||
b2_rect, b2_lbl = batch_box(L_OX, s_y_gpu, s_y_last6, 8, 4, RED, "Batch 2")
|
||||
self.play(Create(b1_rect), Write(b1_lbl))
|
||||
self.play(Create(b2_rect), Write(b2_lbl))
|
||||
self.wait(0.8)
|
||||
|
||||
# ── Right: Continuous Batching ──
|
||||
c_title = Text("Continuous Batching", font_size=26, color=GREEN)
|
||||
c_title.move_to([R_OX + PANEL_W / 2, GY + 0.65, 0])
|
||||
c_note = Text(
|
||||
"no waiting · no padding · GPU never idle",
|
||||
font_size=13, color=GREEN,
|
||||
).move_to([R_OX + PANEL_W / 2, -1.6, 0])
|
||||
self.play(Write(c_title))
|
||||
self.wait(0.25)
|
||||
|
||||
ct_axis = taxis(R_OX, GY)
|
||||
self.play(Create(ct_axis))
|
||||
|
||||
c_y_gpu = GY - ROW
|
||||
cgpu_l = Text("GPU", font_size=14, color=WHITE)
|
||||
cgpu_l.move_to([R_OX - 0.55, c_y_gpu, 0])
|
||||
self.play(Write(cgpu_l))
|
||||
|
||||
c_gpu = gbox(R_OX, c_y_gpu, 0, 12, GREEN, 0.75)
|
||||
self.play(GrowFromEdge(c_gpu, LEFT), run_time=0.5)
|
||||
|
||||
c_reqs = [
|
||||
("A", ORANGE, 0, 4),
|
||||
("B", BLUE, 1, 4),
|
||||
("C", PINK, 2, 4),
|
||||
("D", ORANGE, 4, 4),
|
||||
("E", BLUE, 6, 4),
|
||||
("F", PINK, 8, 4),
|
||||
]
|
||||
c_bars = []
|
||||
c_n_reqs = len(c_reqs)
|
||||
for i, (name, col, start, span) in enumerate(c_reqs):
|
||||
y = c_y_gpu - (i + 1) * ROW
|
||||
lbl = Text(f"Req {name}", font_size=12, color=col)
|
||||
lbl.move_to([R_OX - 0.55, y, 0])
|
||||
bar_rect = gbox(R_OX, y, start, span, col, 0.60)
|
||||
c_bars.extend([lbl, bar_rect])
|
||||
self.play(FadeIn(lbl), GrowFromEdge(bar_rect, LEFT), run_time=0.09)
|
||||
self.wait(0.3)
|
||||
|
||||
# ---- title ----
|
||||
c_y_last = c_y_gpu - c_n_reqs * ROW
|
||||
c_box_rect, c_box_lbl = batch_box(R_OX, c_y_gpu, c_y_last, 0, 12, GREEN, "Always Serving")
|
||||
self.play(Create(c_box_rect), Write(c_box_lbl))
|
||||
self.wait(1.0)
|
||||
|
||||
s_count = Text("6 reqs · 2 batches · GPU idle gaps",
|
||||
font_size=16, color=RED) \
|
||||
.next_to(s_gpu_batch1, DOWN, buff=1.0).align_to(s_gpu_batch1, LEFT)
|
||||
c_count = Text("6 reqs · continuous · GPU never idle",
|
||||
font_size=16, color=GREEN) \
|
||||
.next_to(c_gpu, DOWN, buff=1.0).align_to(c_gpu, LEFT)
|
||||
self.play(Write(s_note), Write(c_note))
|
||||
self.wait(0.3)
|
||||
self.play(Write(s_count), Write(c_count))
|
||||
self.wait(2.5)
|
||||
self.play(FadeOut(s_count), FadeOut(c_count))
|
||||
|
||||
gantt_mobs = [
|
||||
title, bar, s_title, s_note, c_title, c_note,
|
||||
gpu_l, cgpu_l, s_idle1, s_idle2, st_axis, ct_axis,
|
||||
*s_gpu_bars, c_gpu, *s_bars, *c_bars,
|
||||
b1_rect, b1_lbl, b2_rect, b2_lbl, c_box_rect, c_box_lbl,
|
||||
]
|
||||
self.play(*[FadeOut(m) for m in gantt_mobs])
|
||||
self.wait(0.2)
|
||||
|
||||
# ═══════════════════════════════════════════════════
|
||||
# 10. Throughput comparison with animated bars
|
||||
# ═══════════════════════════════════════════════════
|
||||
compare_title = Text("Throughput Comparison", font_size=30, color=BLUE)
|
||||
self.play(Write(compare_title))
|
||||
self.wait(0.2)
|
||||
self.play(compare_title.animate.to_edge(UP).scale(0.55))
|
||||
self.wait(0.2)
|
||||
|
||||
# ---- bar config ----
|
||||
bar_max_w = 5.0
|
||||
bar_h = 0.55
|
||||
row_gap = 0.8
|
||||
|
||||
# ratio: static = baseline (1.0), continuous = 3.4x
|
||||
ratio = 1.0 / 3.4
|
||||
|
||||
# ---- Static Batching row ----
|
||||
s_label = Text("Static Batching", font_size=24, color=RED)
|
||||
s_rect = Rectangle(width=bar_max_w, height=bar_h, color=RED, stroke_width=1.5)
|
||||
s_bar = Rectangle(width=bar_max_w * ratio, height=bar_h,
|
||||
color=RED, fill_opacity=0.55, stroke_width=0)
|
||||
s_bar_rect = Rectangle(
|
||||
width=bar_max_w * ratio, height=bar_h,
|
||||
color=RED, fill_opacity=0.55, stroke_width=0,
|
||||
)
|
||||
s_num = Text("1.0x", font_size=24, color=RED)
|
||||
|
||||
# ---- Continuous Batching row ----
|
||||
c_label = Text("Continuous Batching", font_size=24, color=GREEN)
|
||||
c_rect = Rectangle(width=bar_max_w, height=bar_h, color=GREEN, stroke_width=1.5)
|
||||
c_bar = Rectangle(width=bar_max_w, height=bar_h,
|
||||
color=GREEN, fill_opacity=0.55, stroke_width=0)
|
||||
c_bar_rect = Rectangle(
|
||||
width=bar_max_w, height=bar_h,
|
||||
color=GREEN, fill_opacity=0.55, stroke_width=0,
|
||||
)
|
||||
c_num = Text("3.4x", font_size=24, color=GREEN)
|
||||
|
||||
# position rects first, then align bars
|
||||
s_rect.move_to(ORIGIN + UP * (row_gap / 2 + bar_h / 2))
|
||||
c_rect.move_to(ORIGIN + DOWN * (row_gap / 2 + bar_h / 2))
|
||||
s_bar.align_to(s_rect, LEFT).align_to(s_rect, UP)
|
||||
c_bar.align_to(c_rect, LEFT).align_to(c_rect, UP)
|
||||
s_bar_rect.align_to(s_rect, LEFT).align_to(s_rect, UP)
|
||||
c_bar_rect.align_to(c_rect, LEFT).align_to(c_rect, UP)
|
||||
|
||||
# labels left, nums right
|
||||
s_label.next_to(s_rect, LEFT, buff=0.4)
|
||||
c_label.next_to(c_rect, LEFT, buff=0.4)
|
||||
s_num.next_to(s_rect, RIGHT, buff=0.4)
|
||||
@@ -312,13 +440,11 @@ class ContinuousBatching(Scene):
|
||||
)
|
||||
self.wait(0.3)
|
||||
|
||||
# grow bars
|
||||
self.play(GrowFromEdge(s_bar, LEFT), rate_func=linear, run_time=0.6)
|
||||
self.play(GrowFromEdge(s_bar_rect, LEFT), rate_func=linear, run_time=0.6)
|
||||
self.wait(0.3)
|
||||
self.play(GrowFromEdge(c_bar, LEFT), rate_func=linear, run_time=0.6)
|
||||
self.play(GrowFromEdge(c_bar_rect, LEFT), rate_func=linear, run_time=0.6)
|
||||
self.wait(0.3)
|
||||
|
||||
# show values
|
||||
self.play(Write(s_num), Write(c_num))
|
||||
self.wait(2.5)
|
||||
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
from manim import *
|
||||
|
||||
Text.set_default(font="Times New Roman")
|
||||
|
||||
|
||||
class CTA(Scene):
|
||||
def construct(self):
|
||||
cta = VGroup(
|
||||
Text("AstrAI", font_size=52, color=BLUE),
|
||||
Text("Single GPU · Open Source · 1B params", font_size=22, color=GRAY),
|
||||
Text("github.com/ViperEkura/AstrAI", font_size=18, color=YELLOW),
|
||||
).arrange(DOWN, buff=0.35)
|
||||
cta.move_to(ORIGIN)
|
||||
self.play(Write(cta))
|
||||
self.wait(2.5)
|
||||
self.play(FadeOut(cta))
|
||||
@@ -0,0 +1,97 @@
|
||||
"""AstrAI promo: Hook — terminal demo + GPU comparison + logo reveal (~8s)."""
|
||||
|
||||
from manim import *
|
||||
|
||||
Text.set_default(font="Times New Roman")
|
||||
|
||||
|
||||
class Hook(Scene):
|
||||
def construct(self):
|
||||
BG = "#0d1117"
|
||||
self.camera.background_color = BG
|
||||
TERM = "#39ff14"
|
||||
|
||||
# ════ 0. Terminal (0.0 – 2.8s) ════
|
||||
tw, th = 10.0, 4.8
|
||||
term_frame = Rectangle(
|
||||
width=tw, height=th, color=WHITE, stroke_width=1.2,
|
||||
fill_opacity=0.03, fill_color=TERM,
|
||||
)
|
||||
term_frame.to_edge(UP, buff=0.2)
|
||||
|
||||
lx = term_frame.get_left()[0] + 0.45
|
||||
ty = term_frame.get_top()[1] - 0.45
|
||||
|
||||
prompt = Text("$ ", font_size=20, color=TERM, font="Consolas")
|
||||
prompt.move_to([lx, ty, 0], aligned_edge=LEFT)
|
||||
|
||||
chat = Text("Hey, write a quicksort for me.",
|
||||
font_size=18, color=WHITE, font="Consolas")
|
||||
chat.next_to(prompt, RIGHT, buff=0.08).align_to(prompt, DOWN)
|
||||
|
||||
resp = Text(">> ", font_size=18, color=GRAY, font="Consolas")
|
||||
resp.next_to(prompt, DOWN, buff=0.55).align_to(prompt, LEFT)
|
||||
|
||||
hint = Text("Here is the implementation:",
|
||||
font_size=14, color=GRAY, font="Consolas")
|
||||
hint.next_to(resp, DOWN, buff=0.20).align_to(prompt, LEFT)
|
||||
|
||||
INDENT = 0.27
|
||||
code_lines = [
|
||||
("def quicksort(arr):", 0),
|
||||
("if len(arr) <= 1:", 1),
|
||||
("return arr", 2),
|
||||
("pivot = arr[len(arr)//2]", 1),
|
||||
("left = [x for x in arr if x < pivot]", 1),
|
||||
("right = [x for x in arr if x >= pivot]", 1),
|
||||
("return quicksort(left) + [pivot] + quicksort(right)", 1),
|
||||
]
|
||||
code_vg = VGroup()
|
||||
prev = hint
|
||||
for text, level in code_lines:
|
||||
cl = Text(text, font_size=13, color=TERM, font="Consolas")
|
||||
cl.next_to(prev, DOWN, buff=0.08).align_to(prompt, LEFT).shift(RIGHT * level * INDENT)
|
||||
code_vg.add(cl)
|
||||
prev = cl
|
||||
|
||||
self.play(FadeIn(term_frame), run_time=0.3)
|
||||
self.play(Write(prompt), run_time=0.1)
|
||||
self.play(AddTextLetterByLetter(chat, time_per_char=0.018))
|
||||
self.play(Write(resp), run_time=0.12)
|
||||
self.play(Write(hint, lag_ratio=1.0), run_time=0.15)
|
||||
self.play(Write(code_vg, lag_ratio=1.0), run_time=1.2)
|
||||
|
||||
self.play(
|
||||
FadeOut(term_frame), FadeOut(prompt), FadeOut(chat),
|
||||
FadeOut(resp), FadeOut(hint), FadeOut(code_vg),
|
||||
run_time=0.3,
|
||||
)
|
||||
|
||||
# ════ 1. Statement (2.8 – 4.0s) ════
|
||||
s1 = Text("LLMs Are Powerful", font_size=44, color=WHITE)
|
||||
s2 = Text("but do they really need dozens of GPUs?",
|
||||
font_size=28, color=GRAY)
|
||||
VGroup(s1, s2).arrange(DOWN, buff=0.25).move_to(UP * 1.0)
|
||||
self.play(Write(s1), run_time=0.35)
|
||||
self.play(Write(s2), run_time=0.3)
|
||||
self.wait(0.25)
|
||||
self.play(FadeOut(s1), FadeOut(s2), run_time=0.25)
|
||||
|
||||
# ════ 2. Logo (4.0 – 7.0s) ════
|
||||
logo = Text("AstrAI", font_size=80, color=BLUE)
|
||||
glow = Text("AstrAI", font_size=80, color=BLUE_A, fill_opacity=0.3)
|
||||
glow.move_to(logo).shift(UP * 0.015 + RIGHT * 0.015)
|
||||
|
||||
info = VGroup(
|
||||
Text("1B Params · Single GPU · Open Source",
|
||||
font_size=22, color=GRAY),
|
||||
Text("github.com/ViperEkura/AstrAI", font_size=17, color=YELLOW),
|
||||
).arrange(DOWN, buff=0.18)
|
||||
everything = VGroup(VGroup(glow, logo), info).arrange(DOWN, buff=0.45)
|
||||
everything.move_to(ORIGIN)
|
||||
|
||||
self.play(Write(logo), run_time=0.45)
|
||||
self.play(Write(glow), run_time=0.08)
|
||||
self.play(Write(info), run_time=0.3)
|
||||
self.wait(1.8)
|
||||
self.play(FadeOut(everything), run_time=0.4)
|
||||
+164
@@ -0,0 +1,164 @@
|
||||
"""AstrAI promo: Paged KV Cache — astrai/inference/cache.py & scheduler.py."""
|
||||
|
||||
from manim import *
|
||||
|
||||
Text.set_default(font="Times New Roman")
|
||||
|
||||
|
||||
class _TaskRow:
|
||||
"""Manages one task's logical-page row: label, blocks, arrows."""
|
||||
|
||||
COLS = [-4.0, -3.2, -2.4, -1.6]
|
||||
|
||||
def __init__(self, scene, label, color, y, pool_pos, pool_y):
|
||||
self.scene = scene
|
||||
self.color = color
|
||||
self.y = y
|
||||
self.pool_pos = pool_pos
|
||||
self.pool_y = pool_y
|
||||
self._next_col = 0
|
||||
self.blocks = VGroup()
|
||||
self.arrows = VGroup()
|
||||
|
||||
lbl = scene._small(label, 11, color, weight=BOLD)
|
||||
lbl.move_to([-5.2, y, 0])
|
||||
self._label = lbl
|
||||
|
||||
self.blocks.add(lbl)
|
||||
|
||||
def arrive(self, *phys_idxs):
|
||||
self.scene.play(Write(self._label))
|
||||
for pid in phys_idxs:
|
||||
self._add(pid, expand=False)
|
||||
|
||||
def expand(self, phys_idx):
|
||||
self._add(phys_idx, expand=True)
|
||||
|
||||
def _add(self, phys_idx, expand):
|
||||
col = self._next_col
|
||||
self._next_col += 1
|
||||
x = self.COLS[col]
|
||||
pos = np.array([x, self.y, 0])
|
||||
|
||||
pb = self.scene._lp_box(pos, str(col), self.color)
|
||||
self.blocks.add(pb)
|
||||
|
||||
arr = Arrow(
|
||||
[x, self.y + 0.19, 0],
|
||||
[self.pool_pos[phys_idx][0], self.pool_y - 0.22, 0],
|
||||
color=self.color, stroke_width=1.5, buff=0.03,
|
||||
max_tip_length_to_length_ratio=0.12,
|
||||
)
|
||||
self.arrows.add(arr)
|
||||
|
||||
if expand:
|
||||
target = pb.copy()
|
||||
pb.scale(0)
|
||||
self.scene.add(pb)
|
||||
self.scene.play(Transform(pb, target), GrowArrow(arr), run_time=0.3)
|
||||
else:
|
||||
self.scene.play(FadeIn(pb, scale=0.5), GrowArrow(arr), run_time=0.12)
|
||||
|
||||
def finish(self):
|
||||
self.scene.play(FadeOut(self.blocks), FadeOut(self.arrows))
|
||||
|
||||
|
||||
class PagedCache(Scene):
|
||||
def _small(self, text, size=10, color=GRAY, **kwargs):
|
||||
return Text(text, font_size=size, color=color, **kwargs)
|
||||
|
||||
def _page_box(self, pos, label, color, sz=0.44):
|
||||
s = Square(side_length=sz, color=color, fill_opacity=0.12, stroke_width=1.6)
|
||||
s.move_to(pos)
|
||||
lbl = self._small(label, 10, color).move_to(pos)
|
||||
return VGroup(s, lbl)
|
||||
|
||||
def _lp_box(self, pos, label, color, sz=0.38):
|
||||
s = RoundedRectangle(width=sz, height=sz, corner_radius=0.06,
|
||||
color=color, fill_opacity=0.22, stroke_width=1.6)
|
||||
s.move_to(pos)
|
||||
lbl = self._small(label, 12, color).move_to(pos)
|
||||
return VGroup(s, lbl)
|
||||
|
||||
def construct(self):
|
||||
title = Text("Paged KV Cache", font_size=20, color=BLUE)
|
||||
title.to_edge(UP, buff=0.15)
|
||||
self.play(Write(title))
|
||||
|
||||
pool_y = 1.45; pool_x0 = -3.8; sp = 0.68
|
||||
pool_pages = []; pool_pos = []
|
||||
for i in range(8):
|
||||
x = pool_x0 + i * sp
|
||||
pos = np.array([x, pool_y, 0])
|
||||
pool_pos.append(pos)
|
||||
pb = self._page_box(pos, str(i), GRAY)
|
||||
pool_pages.append(pb)
|
||||
self.play(FadeIn(pb, scale=0.5), run_time=0.04)
|
||||
self.wait(0.1)
|
||||
|
||||
plbl = self._small("page frames [0..7]", 9, GRAY)
|
||||
plbl.next_to(pool_pages[0][0], DOWN, buff=0.25).shift(LEFT * 0.3)
|
||||
self.play(Write(plbl))
|
||||
|
||||
mask = self._small("free: 11111111", 10, GRAY)
|
||||
mask.next_to(plbl, DOWN, buff=0.1, aligned_edge=LEFT)
|
||||
self.play(Write(mask))
|
||||
|
||||
def alloc(idx, color):
|
||||
pg = pool_pages[idx][0]
|
||||
self.play(pg.animate.set_fill(color, opacity=0.35), run_time=0.1)
|
||||
flash = SurroundingRectangle(pool_pages[idx], color=color, buff=0.04)
|
||||
self.play(Create(flash), run_time=0.05)
|
||||
self.play(FadeOut(flash), run_time=0.04)
|
||||
|
||||
def set_mask(bits):
|
||||
m2 = self._small(f"free: {bits}", 10, GRAY)
|
||||
m2.next_to(plbl, DOWN, buff=0.1, aligned_edge=LEFT)
|
||||
self.play(Transform(mask, m2))
|
||||
|
||||
a_y = 0.25; b_y = -0.45; c_y = -1.15
|
||||
|
||||
# ── A arrives ──
|
||||
alloc(0, GREEN); alloc(1, GREEN); set_mask("11111100")
|
||||
a = _TaskRow(self, "A", GREEN, a_y, pool_pos, pool_y)
|
||||
a.arrive(0, 1)
|
||||
|
||||
# ── A expands 1 ──
|
||||
alloc(4, GREEN); set_mask("11101100")
|
||||
a.expand(4)
|
||||
|
||||
# ── B arrives ──
|
||||
alloc(2, ORANGE); alloc(3, ORANGE); set_mask("11100000")
|
||||
b = _TaskRow(self, "B", ORANGE, b_y, pool_pos, pool_y)
|
||||
b.arrive(2, 3)
|
||||
|
||||
# ── C arrives ──
|
||||
alloc(5, BLUE); alloc(6, BLUE); set_mask("10000000")
|
||||
c = _TaskRow(self, "C", BLUE, c_y, pool_pos, pool_y)
|
||||
c.arrive(5, 6)
|
||||
|
||||
# ── A expands 2 ──
|
||||
alloc(7, GREEN); set_mask("00000000")
|
||||
a.expand(7)
|
||||
|
||||
# ── A finishes ──
|
||||
a.finish()
|
||||
for idx in [0, 1, 4, 7]:
|
||||
pg = pool_pages[idx][0]
|
||||
self.play(pg.animate.set_fill(GRAY, opacity=0.12), run_time=0.08)
|
||||
flash = SurroundingRectangle(pool_pages[idx], color=YELLOW, buff=0.04)
|
||||
self.play(Create(flash), run_time=0.06)
|
||||
self.play(FadeOut(flash), run_time=0.04)
|
||||
set_mask("10010011")
|
||||
|
||||
# ── B expands (reuse) ──
|
||||
alloc(0, ORANGE); set_mask("10010010")
|
||||
b.expand(0)
|
||||
|
||||
self.wait(0.5)
|
||||
s = Text("Page-table-indirected, O(1) alloc/free, on-demand growth",
|
||||
font_size=12, color=GREEN)
|
||||
s.move_to([-3.0, -2.0, 0])
|
||||
self.play(Write(s))
|
||||
self.wait(2)
|
||||
self.play(*[FadeOut(m) for m in self.mobjects])
|
||||
-117
@@ -1,117 +0,0 @@
|
||||
"""AstrAI promo: Prefix Cache animation (Radix tree with branches)."""
|
||||
|
||||
from manim import *
|
||||
|
||||
|
||||
class PrefixCache(Scene):
|
||||
"""Animates the radix-tree prefix cache with multiple distinct branches."""
|
||||
|
||||
def _add_node(self, parent_pos, label, color, dx, dy):
|
||||
pos = parent_pos + np.array([dx, dy, 0])
|
||||
dot = Dot(point=pos, color=color, radius=0.1)
|
||||
txt = Text(label, font_size=13, color=color)
|
||||
txt.next_to(dot, UP, buff=0.1)
|
||||
grp = VGroup(dot, txt)
|
||||
edge = Line(parent_pos, pos, color=GRAY, stroke_width=1.5)
|
||||
return grp, edge, pos
|
||||
|
||||
def _add_leaf(self, parent_pos, color, tag):
|
||||
leaf = Square(side_length=0.25, color=color, fill_opacity=0.4)
|
||||
leaf.move_to(parent_pos + DOWN * 0.7)
|
||||
edge = Line(parent_pos, leaf.get_top(), color=color, stroke_width=1.5)
|
||||
lbl = Text(tag, font_size=10, color=color).next_to(leaf, DOWN, buff=0.1)
|
||||
return VGroup(leaf, edge, lbl)
|
||||
|
||||
def construct(self):
|
||||
title = Text("Prefix Cache", font_size=48, color=BLUE)
|
||||
self.play(Write(title))
|
||||
self.wait(0.2)
|
||||
self.play(title.animate.to_edge(UP).scale(0.6))
|
||||
|
||||
# Root at top-left, tree stays visible throughout
|
||||
root_pos = np.array([-4.5, 2.0, 0])
|
||||
root = Circle(radius=0.25, color=BLUE, fill_opacity=0.2)
|
||||
root.move_to(root_pos)
|
||||
root_lbl = Text("root", font_size=10, color=GRAY).move_to(root)
|
||||
root_grp = VGroup(root, root_lbl)
|
||||
self.play(FadeIn(root_grp, scale=0.5), run_time=0.3)
|
||||
|
||||
# Labels accumulate on the right side
|
||||
right_x = 3.5
|
||||
label_y = 2.5
|
||||
label_step = 0.5
|
||||
|
||||
def show_label(text, color):
|
||||
nonlocal label_y
|
||||
lbl = Text(text, font_size=14, color=color)
|
||||
lbl.move_to([right_x, label_y, 0])
|
||||
label_y -= label_step
|
||||
self.play(Write(lbl))
|
||||
return lbl
|
||||
|
||||
# ── R1: A → B → C ──
|
||||
r1_lbl = show_label('R1: "A B C"', GREEN)
|
||||
|
||||
a_grp, a_edge, a_pos = self._add_node(root_pos, "A", GREEN, 0.6, -0.9)
|
||||
self.play(Create(a_edge), FadeIn(a_grp, scale=0.5), run_time=0.2)
|
||||
b_grp, b_edge, b_pos = self._add_node(a_pos, "B", GREEN, 0.6, -0.9)
|
||||
self.play(Create(b_edge), FadeIn(b_grp, scale=0.5), run_time=0.2)
|
||||
c_grp, c_edge, c_pos = self._add_node(b_pos, "C", GREEN, 0.6, -0.9)
|
||||
self.play(Create(c_edge), FadeIn(c_grp, scale=0.5), run_time=0.2)
|
||||
self.play(FadeIn(self._add_leaf(c_pos, GREEN, "slot 0"), scale=0.8), run_time=0.3)
|
||||
self.wait(0.3)
|
||||
|
||||
# ── R2: shares A B, branches D E ──
|
||||
r2_lbl = show_label('R2: "A B D E"', ORANGE)
|
||||
|
||||
for g in [a_grp, b_grp]:
|
||||
flash = SurroundingRectangle(g, color=YELLOW, buff=0.12)
|
||||
self.play(Create(flash), run_time=0.1)
|
||||
self.play(FadeOut(flash), run_time=0.08)
|
||||
|
||||
d_grp, d_edge, d_pos = self._add_node(b_pos, "D", ORANGE, -0.6, -0.9)
|
||||
self.play(Create(d_edge), FadeIn(d_grp, scale=0.5), run_time=0.2)
|
||||
e_grp, e_edge, e_pos = self._add_node(d_pos, "E", ORANGE, -0.6, -0.9)
|
||||
self.play(Create(e_edge), FadeIn(e_grp, scale=0.5), run_time=0.2)
|
||||
self.play(FadeIn(self._add_leaf(e_pos, ORANGE, "slot 1"), scale=0.8), run_time=0.3)
|
||||
self.wait(0.3)
|
||||
|
||||
# ── R3: shares A B, single F ──
|
||||
r3_lbl = show_label('R3: "A B F"', PINK)
|
||||
|
||||
f_grp, f_edge, f_pos = self._add_node(b_pos, "F", PINK, 0.0, -1.2)
|
||||
self.play(Create(f_edge), FadeIn(f_grp, scale=0.5), run_time=0.2)
|
||||
self.play(FadeIn(self._add_leaf(f_pos, PINK, "slot 2"), scale=0.8), run_time=0.3)
|
||||
self.wait(0.3)
|
||||
|
||||
# ── R4: new prefix from root ──
|
||||
r4_lbl = show_label('R4: "X Y"', TEAL)
|
||||
|
||||
x_grp, x_edge, x_pos = self._add_node(root_pos, "X", TEAL, -1.0, -0.9)
|
||||
self.play(Create(x_edge), FadeIn(x_grp, scale=0.5), run_time=0.2)
|
||||
y_grp, y_edge, y_pos = self._add_node(x_pos, "Y", TEAL, -0.6, -0.9)
|
||||
self.play(Create(y_edge), FadeIn(y_grp, scale=0.5), run_time=0.2)
|
||||
self.play(FadeIn(self._add_leaf(y_pos, TEAL, "slot 3"), scale=0.8), run_time=0.3)
|
||||
self.wait(0.5)
|
||||
|
||||
# ── highlight shared prefix (tree stays) ──
|
||||
reuse_box = SurroundingRectangle(VGroup(a_grp, b_grp), color=YELLOW, buff=0.15)
|
||||
reuse_note = Text(
|
||||
'Prefix "A B" shared\nby 3 requests — 0 copy',
|
||||
font_size=16,
|
||||
color=YELLOW,
|
||||
)
|
||||
reuse_note.next_to(reuse_box, LEFT, buff=1.0)
|
||||
self.play(Create(reuse_box), Write(reuse_note))
|
||||
self.wait(2)
|
||||
self.play(FadeOut(reuse_box), FadeOut(reuse_note))
|
||||
|
||||
# ── summary below tree (tree stays visible) ──
|
||||
summary = VGroup(
|
||||
Text("KV cache reuse across requests", font_size=26, color=GREEN),
|
||||
Text("First-token latency: up to 50% reduction", font_size=18, color=GRAY),
|
||||
).arrange(DOWN, buff=0.2)
|
||||
summary.to_edge(DOWN, buff=0.5)
|
||||
self.play(Write(summary))
|
||||
self.wait(2)
|
||||
self.play(FadeOut(summary), FadeOut(root_grp), FadeOut(title))
|
||||
+16
-6
@@ -1,16 +1,19 @@
|
||||
"""Render all promo scenes with Manim."""
|
||||
"""Render all promo scenes with Manim in parallel."""
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from pathlib import Path
|
||||
|
||||
ROOT = Path(__file__).parent
|
||||
|
||||
SCENES = [
|
||||
("hook.py", "Hook"),
|
||||
("transformer.py", "Transformer"),
|
||||
("architecture.py", "Architecture"),
|
||||
("continuous_batching.py", "ContinuousBatching"),
|
||||
("prefix_cache.py", "PrefixCache"),
|
||||
("paged_cache.py", "PrefixCache"),
|
||||
("cta.py", "CTA"),
|
||||
]
|
||||
|
||||
|
||||
@@ -27,15 +30,22 @@ def render(file_name, scene_name, quality="-qh"):
|
||||
"--media_dir",
|
||||
str(media_dir),
|
||||
]
|
||||
print(f"Rendering {scene_name}...")
|
||||
print(f"[{scene_name}] Rendering...")
|
||||
subprocess.run(cmd, check=True)
|
||||
print(f" Done → {media_dir / 'videos' / scene_name.lower()}.mp4")
|
||||
print(f"[{scene_name}] Done → {media_dir / 'videos' / scene_name.lower()}.mp4")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
quality = "-qh" # 1080p; use -l for draft, -4k for ultra
|
||||
if len(sys.argv) > 1:
|
||||
quality = sys.argv[1]
|
||||
for f, s in SCENES:
|
||||
render(f, s, quality)
|
||||
|
||||
max_workers = len(SCENES)
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as pool:
|
||||
fut = {pool.submit(render, f, s, quality): s for f, s in SCENES}
|
||||
for f in as_completed(fut):
|
||||
exc = f.exception()
|
||||
if exc:
|
||||
print(f"[{fut[f]}] Failed: {exc}")
|
||||
|
||||
print("All scenes rendered.")
|
||||
|
||||
+405
-56
@@ -6,18 +6,21 @@ Shows the Grouped-Query Attention (GQA) mechanism with orthogonal data-flow line
|
||||
|
||||
from manim import *
|
||||
import numpy as np
|
||||
import math
|
||||
|
||||
Text.set_default(font="Times New Roman")
|
||||
|
||||
|
||||
class Transformer(Scene):
|
||||
"""Animates the GQA attention mechanism with orthogonal connection lines."""
|
||||
|
||||
def construct(self):
|
||||
title = Text("Grouped-Query Attention (GQA)", font_size=42, color=BLUE)
|
||||
title = Text("Grouped-Query Attention (GQA)", font_size=34, color=BLUE)
|
||||
title.to_edge(UP, buff=0.35)
|
||||
self.play(Write(title))
|
||||
|
||||
# ── Helper: box ──
|
||||
def mk(name, color, w=2.6, h=0.72, fs=10):
|
||||
def mk(name, color, w=3.0, h=0.85, fs=13):
|
||||
box = Rectangle(
|
||||
width=w, height=h, color=color, fill_opacity=0.12, stroke_width=1.5
|
||||
)
|
||||
@@ -25,32 +28,32 @@ class Transformer(Scene):
|
||||
return VGroup(box, lbl)
|
||||
|
||||
# ── Layout ──
|
||||
inp = Text("x (hidden states)", font_size=15, color=GRAY)
|
||||
inp.move_to(UP * 2.8)
|
||||
inp = Text("x (hidden states)", font_size=20, color=GRAY)
|
||||
inp.move_to(UP * 2.5)
|
||||
|
||||
y1 = 1.5
|
||||
y1 = 1.6
|
||||
q_grp = mk("Q Projection\n1536 → 24×64", YELLOW)
|
||||
k_grp = mk("K Projection\n1536 → 4×64", YELLOW)
|
||||
v_grp = mk("V Projection\n1536 → 4×64", YELLOW)
|
||||
q_grp.move_to(LEFT * 3.0 + UP * y1)
|
||||
q_grp.move_to(LEFT * 3.6 + UP * y1)
|
||||
k_grp.move_to(UP * y1)
|
||||
v_grp.move_to(RIGHT * 3.0 + UP * y1)
|
||||
v_grp.move_to(RIGHT * 3.6 + UP * y1)
|
||||
|
||||
y2 = 0.0
|
||||
repeat_grp = mk("Repeat KV\n4 heads → 24 heads", GREEN, 2.4, 0.68, 10)
|
||||
y2 = 0.4
|
||||
repeat_grp = mk("Repeat KV\n4 heads → 24 heads", GREEN, 2.8, 0.80)
|
||||
repeat_grp.move_to(UP * y2)
|
||||
|
||||
y3 = -1.6
|
||||
y3 = -1.0
|
||||
sdpa_grp = mk(
|
||||
"Scaled Dot-Product\nAttention Q·Kᵀ/√d", BLUE, 2.8, 0.74, 10
|
||||
"Scaled Dot-Product\nAttention Q·K^T/√d", BLUE, 3.2, 0.85,
|
||||
)
|
||||
sdpa_grp.move_to(UP * y3)
|
||||
|
||||
y4 = -3.0
|
||||
o_grp = mk("O Projection\n1536 → 1536", PURPLE, 2.2, 0.68, 10)
|
||||
y4 = -2.2
|
||||
o_grp = mk("O Projection\n1536 → 1536", PURPLE, 2.6, 0.80)
|
||||
o_grp.move_to(UP * y4)
|
||||
|
||||
out = Text("x' (hidden states)", font_size=15, color=GRAY)
|
||||
out = Text("x' (hidden states)", font_size=20, color=GRAY)
|
||||
out.next_to(o_grp, DOWN, buff=0.4)
|
||||
|
||||
# ── Animate boxes ──
|
||||
@@ -60,7 +63,7 @@ class Transformer(Scene):
|
||||
self.play(FadeIn(g, shift=UP * 0.1), run_time=0.2)
|
||||
|
||||
# ── Input trunk → branch → Q/K/V (enter from directly above) ──
|
||||
trunk_bottom = np.array([0, q_grp.get_top()[1] + 0.35, 0])
|
||||
trunk_bottom = np.array([0, q_grp.get_top()[1] + 0.2, 0])
|
||||
trunk = Line(inp.get_bottom(), trunk_bottom, color=GRAY, stroke_width=1.5)
|
||||
self.play(Create(trunk), run_time=0.15)
|
||||
|
||||
@@ -166,59 +169,405 @@ class Transformer(Scene):
|
||||
VGroup(q_grp, k_grp, v_grp), color=YELLOW, buff=0.2
|
||||
)
|
||||
gqa_t = Text(
|
||||
"GQA 6:1 — 24 Q-heads → 4 KV-heads\nKV cache reduced by 83%",
|
||||
font_size=13, color=YELLOW,
|
||||
"GQA 6:1\n24 Q-heads → 4 KV-heads\nKV cache -83%",
|
||||
font_size=11, color=YELLOW,
|
||||
)
|
||||
gqa_t.next_to(gqa_h, RIGHT, buff=0.5)
|
||||
gqa_t.next_to(v_grp, DOWN, buff=0.4).shift(RIGHT * 0.4)
|
||||
self.play(Create(gqa_h), Write(gqa_t))
|
||||
self.wait(1.8)
|
||||
|
||||
# ── Repeat KV highlight ──
|
||||
kv_h = SurroundingRectangle(
|
||||
VGroup(k_grp, v_grp), color=GREEN, buff=0.12
|
||||
)
|
||||
kv_t = Text(
|
||||
"repeat_kv(): broadcast\n4 heads → 24 heads",
|
||||
font_size=12, color=GREEN,
|
||||
)
|
||||
kv_t.next_to(kv_h, RIGHT, buff=0.5)
|
||||
self.play(Create(kv_h), Write(kv_t))
|
||||
self.wait(1.5)
|
||||
self.play(FadeOut(gqa_h), FadeOut(gqa_t))
|
||||
|
||||
# ── Fade all ──
|
||||
self.play(
|
||||
*[FadeOut(g) for g in all_boxes],
|
||||
FadeOut(all_lines),
|
||||
FadeOut(kv_h), FadeOut(kv_t),
|
||||
FadeOut(gqa_h), FadeOut(gqa_t),
|
||||
FadeOut(inp), FadeOut(out), FadeOut(title),
|
||||
)
|
||||
|
||||
# ── Specs card ──
|
||||
st = Text("Model Specifications", font_size=36, color=BLUE)
|
||||
st.to_edge(UP, buff=0.5)
|
||||
rows_data = [
|
||||
("Parameters", "~1.0B"),
|
||||
("Layers", "24 × DecoderBlock"),
|
||||
("Hidden Dim", "1536"),
|
||||
("Q Heads / KV Heads", "24 / 4 (GQA, 6:1)"),
|
||||
("Head Dim", "64"),
|
||||
("FFN Dim", "4608 (SwiGLU)"),
|
||||
("Max Length", "2048"),
|
||||
("Precision", "bfloat16"),
|
||||
# ═══════════════════════════════════════════════════
|
||||
# 12. Scaled Dot-Product Attention — full formula + breakdown
|
||||
# ═══════════════════════════════════════════════════
|
||||
qkv_title = Text("Scaled Dot-Product Attention", font_size=34, color=BLUE)
|
||||
qkv_title.to_edge(UP, buff=0.35)
|
||||
self.play(Write(qkv_title))
|
||||
|
||||
# Full formula first, stays on screen
|
||||
full_eq = MathTex(
|
||||
r"\operatorname{Attention}(Q,K,V)=\operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)\!V",
|
||||
font_size=36, color=WHITE,
|
||||
)
|
||||
full_eq.next_to(qkv_title, DOWN, buff=0.5)
|
||||
self.play(Write(full_eq))
|
||||
self.wait(0.6)
|
||||
|
||||
# Q / K / V — brief meanings
|
||||
q_txt = Text("Q = Query", font_size=18, color=YELLOW)
|
||||
k_txt = Text("K = Key", font_size=18, color=ORANGE)
|
||||
v_txt = Text("V = Value", font_size=18, color=GREEN)
|
||||
qkv_labels = VGroup(q_txt, k_txt, v_txt).arrange(RIGHT, buff=1.5)
|
||||
qkv_labels.next_to(full_eq, DOWN, buff=0.5)
|
||||
self.play(Write(qkv_labels))
|
||||
|
||||
q_desc = Text("\"what am I looking for?\"", font_size=11, color=YELLOW) \
|
||||
.next_to(q_txt, DOWN, buff=0.10)
|
||||
k_desc = Text("\"what do I have?\"", font_size=11, color=ORANGE) \
|
||||
.next_to(k_txt, DOWN, buff=0.10)
|
||||
v_desc = Text("\"what do I contribute?\"", font_size=11, color=GREEN) \
|
||||
.next_to(v_txt, DOWN, buff=0.10)
|
||||
self.play(Write(q_desc), Write(k_desc), Write(v_desc))
|
||||
self.wait(2.0)
|
||||
self.play(FadeOut(qkv_labels), FadeOut(q_desc), FadeOut(k_desc), FadeOut(v_desc))
|
||||
|
||||
# Step-by-step decomposition — full formula stays visible above
|
||||
steps = [
|
||||
(MathTex(r"\text{(1) } S = QK^\top", font_size=26, color=YELLOW),
|
||||
Text("score matrix — pairwise token similarity", font_size=12, color=GRAY)),
|
||||
(MathTex(r"\text{(2) } S \mathbin{/} \sqrt{d_k}", font_size=26, color=ORANGE),
|
||||
Text("scale — prevents gradient explosion", font_size=12, color=GRAY)),
|
||||
(MathTex(r"\text{(3) } \operatorname{softmax}\!\left(S \mathbin{/} \sqrt{d_k}\right)", font_size=26, color=GREEN),
|
||||
Text("normalize — each row sums to 1 (probability)", font_size=12, color=GRAY)),
|
||||
(MathTex(r"\text{(4) } \operatorname{softmax}\!\left(S \mathbin{/} \sqrt{d_k}\right) \cdot V", font_size=26, color=BLUE),
|
||||
Text("weighted sum — aggregate values by attention", font_size=12, color=GRAY)),
|
||||
]
|
||||
table = VGroup()
|
||||
for label, value in rows_data:
|
||||
row = VGroup(
|
||||
Text(label + ":", font_size=15, color=GRAY),
|
||||
Text(value, font_size=15, color=WHITE),
|
||||
).arrange(RIGHT, buff=0.4, aligned_edge=LEFT)
|
||||
table.add(row)
|
||||
table.arrange(DOWN, buff=0.1, aligned_edge=LEFT)
|
||||
table.next_to(st, DOWN, buff=0.4)
|
||||
self.play(Write(st), Write(table))
|
||||
self.wait(2)
|
||||
self.play(FadeOut(st), FadeOut(table))
|
||||
|
||||
step_group = VGroup()
|
||||
steps_mobj = VGroup()
|
||||
for eq, desc in steps:
|
||||
sg = VGroup(eq, desc).arrange(DOWN, buff=0.06)
|
||||
step_group.add(sg)
|
||||
steps_mobj.add(eq)
|
||||
step_group.arrange(DOWN, buff=0.22, aligned_edge=LEFT)
|
||||
step_group.next_to(full_eq, DOWN, buff=0.6)
|
||||
|
||||
for sg in step_group:
|
||||
self.play(Write(sg), run_time=0.3)
|
||||
self.wait(2.5)
|
||||
self.play(FadeOut(step_group), FadeOut(full_eq))
|
||||
|
||||
# ═══════════════════════════════════════════════════
|
||||
# 13. Attention score heatmap — concrete example
|
||||
# ═══════════════════════════════════════════════════
|
||||
hm_title = Text("Attention Score Heatmap", font_size=34, color=BLUE)
|
||||
hm_title.to_edge(UP, buff=0.35)
|
||||
hm_sub = Text("\"The cat sat on the mat\" — causal, per-token attention weights",
|
||||
font_size=14, color=GRAY).next_to(hm_title, DOWN, buff=0.12)
|
||||
self.play(FadeOut(qkv_title), Write(hm_title), Write(hm_sub))
|
||||
|
||||
tokens = ["<s>", "The", "cat", "sat", "on", "the", "mat"]
|
||||
n = len(tokens)
|
||||
cell_size = 0.65
|
||||
gap = 0.05
|
||||
grid_high = n * cell_size + (n - 1) * gap
|
||||
grid_left = -grid_high / 2
|
||||
grid_top = 1.7
|
||||
|
||||
# pre-mask raw scores (QK^T / sqrt(d_k)) — random-varied, distance-biased
|
||||
pre_scores = [
|
||||
[2.8, 1.5, 0.3, 0.1, 0.0, 0.0, 0.0],
|
||||
[1.2, 3.5, 1.8, 0.5, 0.2, 0.1, 0.0],
|
||||
[0.4, 2.0, 3.0, 1.5, 0.6, 0.2, 0.1],
|
||||
[0.1, 0.6, 2.5, 2.8, 1.2, 0.4, 0.1],
|
||||
[0.0, 0.2, 0.8, 2.0, 2.5, 1.5, 0.3],
|
||||
[0.0, 0.1, 0.3, 0.9, 1.8, 2.5, 1.2],
|
||||
[0.0, 0.0, 0.1, 0.4, 0.8, 1.5, 3.0],
|
||||
]
|
||||
|
||||
# compute post-softmax weights with causal mask (j > i → -inf)
|
||||
post_weights = []
|
||||
for i in range(n):
|
||||
row = pre_scores[i]
|
||||
masked = [-float('inf') if j > i else row[j] for j in range(n)]
|
||||
exps = [math.exp(v) for v in masked]
|
||||
exp_sum = sum(exps)
|
||||
post_weights.append([e / exp_sum for e in exps])
|
||||
|
||||
flat_pre = [w for row in pre_scores for w in row]
|
||||
pre_min, pre_max = min(flat_pre), max(flat_pre)
|
||||
flat_post = [w for row in post_weights for w in row]
|
||||
post_min, post_max = min(flat_post), max(flat_post)
|
||||
cells = VGroup()
|
||||
masked_cells = VGroup()
|
||||
for i in range(n):
|
||||
for j in range(n):
|
||||
pw = pre_scores[i][j]
|
||||
pw_normed = (pw - pre_min) / (pre_max - pre_min)
|
||||
pw_color = interpolate_color(BLUE, RED, pw_normed)
|
||||
sq = Square(
|
||||
side_length=cell_size, fill_color=pw_color,
|
||||
fill_opacity=0.75, stroke_width=0.5,
|
||||
stroke_color=GRAY,
|
||||
)
|
||||
x = grid_left + j * (cell_size + gap) + cell_size / 2
|
||||
y = grid_top - i * (cell_size + gap) - cell_size / 2
|
||||
sq.move_to([x, y, 0])
|
||||
cells.add(sq)
|
||||
if j > i:
|
||||
masked_cells.add(sq)
|
||||
self.play(FadeIn(sq, scale=0.6), run_time=0.015)
|
||||
|
||||
# row labels (query) on the left
|
||||
row_lbls = VGroup()
|
||||
for i, tok in enumerate(tokens):
|
||||
lbl = Text(tok, font_size=14, color=GRAY)
|
||||
y = grid_top - i * (cell_size + gap) - cell_size / 2
|
||||
lbl.next_to([grid_left - 0.15, y, 0], LEFT, buff=0.08)
|
||||
row_lbls.add(lbl)
|
||||
q_label = Text("Q", font_size=11, color=WHITE)
|
||||
q_label.move_to(row_lbls[0].get_left() + LEFT * 0.3).shift(UP * 0.15)
|
||||
self.play(*[Write(l) for l in row_lbls], Write(q_label))
|
||||
|
||||
# column labels (key) on top
|
||||
col_lbls = VGroup()
|
||||
for j, tok in enumerate(tokens):
|
||||
lbl = Text(tok, font_size=10, color=GRAY).rotate(PI / 6)
|
||||
x = grid_left + j * (cell_size + gap) + cell_size / 2
|
||||
lbl.next_to([x, grid_top + 0.06, 0], UP, buff=0.04)
|
||||
col_lbls.add(lbl)
|
||||
k_label = Text("K", font_size=11, color=WHITE)
|
||||
k_label.next_to(col_lbls[0], UP, buff=0.06)
|
||||
self.play(*[Write(l) for l in col_lbls], Write(k_label))
|
||||
self.wait(1.0)
|
||||
|
||||
# causal mask + softmax — zero out future tokens, recompute weights
|
||||
causal_txt = Text("causal mask + softmax\n(future tokens → 0)", font_size=11, color=RED) \
|
||||
.next_to(cells[n - 1], UP, buff=0.25).align_to(cells[n - 1], RIGHT)
|
||||
anims = [sq.animate.set_fill(DARK_GRAY, 0.15) for sq in masked_cells]
|
||||
for i in range(n):
|
||||
for j in range(n):
|
||||
if j <= i:
|
||||
idx = i * n + j
|
||||
aw = post_weights[i][j]
|
||||
aw_normed = (aw - post_min) / (post_max - post_min)
|
||||
aw_color = interpolate_color(BLUE, RED, aw_normed)
|
||||
anims.append(cells[idx].animate.set_fill(aw_color, 0.75))
|
||||
self.play(*anims, Write(causal_txt))
|
||||
self.wait(1.2)
|
||||
self.play(FadeOut(causal_txt))
|
||||
|
||||
# highlight key patterns
|
||||
h1 = SurroundingRectangle(cells[2 * n + 1], color=ORANGE, stroke_width=2, buff=0.04)
|
||||
h2 = SurroundingRectangle(cells[3 * n + 2], color=ORANGE, stroke_width=2, buff=0.04)
|
||||
h3 = SurroundingRectangle(cells[4 * n + 3], color=ORANGE, stroke_width=2, buff=0.04)
|
||||
h4 = SurroundingRectangle(cells[5 * n + 4], color=ORANGE, stroke_width=2, buff=0.04)
|
||||
h5 = SurroundingRectangle(cells[6 * n + 5], color=ORANGE, stroke_width=2, buff=0.04)
|
||||
hl_text = Text("previous token attends to next\n(causal sequence learning)", font_size=11, color=ORANGE) \
|
||||
.next_to(cells[(n - 1) * n + (n - 1)], RIGHT, buff=0.8)
|
||||
self.play(Create(h1), Create(h2), Create(h3), Create(h4), Create(h5), Write(hl_text))
|
||||
self.wait(2.0)
|
||||
self.play(FadeOut(h1), FadeOut(h2), FadeOut(h3), FadeOut(h4), FadeOut(h5), FadeOut(hl_text))
|
||||
|
||||
# fade all heatmap
|
||||
self.play(
|
||||
FadeOut(hm_title), FadeOut(hm_sub),
|
||||
FadeOut(cells), FadeOut(row_lbls), FadeOut(col_lbls),
|
||||
FadeOut(q_label), FadeOut(k_label),
|
||||
)
|
||||
|
||||
# ═══════════════════════════════════════════════════
|
||||
# Auto-regressive Generation Demo (v2: full I/O pipeline)
|
||||
def tok_card(text, fill=DARK_BLUE, stroke=GRAY):
|
||||
t = Text(text, font_size=14, color=WHITE)
|
||||
box = RoundedRectangle(
|
||||
width=t.width + 0.3, height=t.height + 0.16,
|
||||
corner_radius=0.06, fill_color=fill, fill_opacity=0.5,
|
||||
stroke_color=stroke, stroke_width=0.5,
|
||||
)
|
||||
t.move_to(box)
|
||||
return VGroup(box, t)
|
||||
|
||||
gen_tokens = ["<s>", "The", "cat", "sat", "on", "the", "mat"]
|
||||
|
||||
gen_title = Text("Auto-regressive Generation", font_size=34, color=BLUE)
|
||||
gen_title.to_edge(UP, buff=0.35)
|
||||
self.play(Write(gen_title))
|
||||
|
||||
# ── Layout constants ──
|
||||
BLK_W = 1.8
|
||||
BLK_H = 0.28
|
||||
TFR_H = 0.55
|
||||
CX_BLK = -1.0
|
||||
Y_EMB = 1.95
|
||||
Y_TFR = 1.10
|
||||
Y_NORM2 = 0.35
|
||||
Y_HEAD = -0.15
|
||||
|
||||
def mkblk(w, h, y, txt, color, fs=10):
|
||||
b = RoundedRectangle(width=w, height=h, corner_radius=0.06,
|
||||
fill_color=DARK_BLUE, fill_opacity=0.25,
|
||||
stroke_color=color, stroke_width=1.5)
|
||||
l = Text(txt, font_size=fs, color=color).move_to(b)
|
||||
return VGroup(b, l).move_to([CX_BLK, y, 0])
|
||||
|
||||
emb_node = mkblk(BLK_W, BLK_H, Y_EMB, "Embedding", YELLOW, 9)
|
||||
tfr_node = mkblk(2.4, TFR_H, Y_TFR, "Transformer Block\n× 24", PURPLE, 12)
|
||||
norm2 = mkblk(BLK_W, BLK_H, Y_NORM2, "RMS Norm", GREEN, 9)
|
||||
head = mkblk(BLK_W, BLK_H, Y_HEAD, "LM Head", RED, 9)
|
||||
pipeline = VGroup(emb_node, tfr_node, norm2, head)
|
||||
|
||||
# Stack effect for ×24 layers behind Transformer block
|
||||
layer_stack = VGroup()
|
||||
for i in range(3, 0, -1):
|
||||
shadow = RoundedRectangle(
|
||||
width=2.4, height=TFR_H, corner_radius=0.06,
|
||||
stroke_color=BLUE_D, stroke_width=1.0, fill_opacity=0,
|
||||
).move_to([CX_BLK + i * 0.04, Y_TFR - i * 0.04, 0])
|
||||
layer_stack.add(shadow)
|
||||
|
||||
# Arrows between blocks
|
||||
arrows = VGroup()
|
||||
for a, b in zip(pipeline[:-1], pipeline[1:]):
|
||||
arr = Arrow(a.get_bottom(), b.get_top(),
|
||||
color=GRAY, stroke_width=1.2, tip_length=0.07)
|
||||
arrows.add(arr)
|
||||
|
||||
# ── KV Cache (right of Transformer) ──
|
||||
KV_X = 1.8
|
||||
kv_size = 0.16
|
||||
kv_gap = 0.04
|
||||
K_Y = Y_TFR + 0.05
|
||||
V_Y = Y_TFR - 0.22
|
||||
cache_lbl = Text("KV Cache", font_size=8, color=GRAY).move_to([KV_X, Y_TFR + TFR_H / 2 + 0.2, 0])
|
||||
k_hdr = Text("K:", font_size=7, color=YELLOW).move_to([KV_X - 0.65, K_Y, 0])
|
||||
v_hdr = Text("V:", font_size=7, color=ORANGE).move_to([KV_X - 0.65, V_Y, 0])
|
||||
|
||||
kv_k = VGroup()
|
||||
kv_v = VGroup()
|
||||
k0 = Square(kv_size, fill_color=YELLOW, fill_opacity=0.4,
|
||||
stroke_color=YELLOW, stroke_width=0.5).move_to([KV_X, K_Y, 0])
|
||||
v0 = Square(kv_size, fill_color=ORANGE, fill_opacity=0.4,
|
||||
stroke_color=ORANGE, stroke_width=0.5).move_to([KV_X, V_Y, 0])
|
||||
kv_k.add(k0); kv_v.add(v0)
|
||||
kv_group = VGroup(cache_lbl, k_hdr, v_hdr, kv_k, kv_v)
|
||||
|
||||
# ── Distribution builder ──
|
||||
def build_dist(probs, y_center, max_w=3.0):
|
||||
bars = VGroup(); lbls = VGroup()
|
||||
bh = 0.18; bg = 0.03; lx = CX_BLK - max_w / 2
|
||||
items = list(probs.items())
|
||||
n = len(items)
|
||||
y_top = y_center + (n * bh + (n - 1) * bg) / 2
|
||||
for i, (tok, pct) in enumerate(items):
|
||||
w = max_w * pct / 100
|
||||
y = y_top - i * (bh + bg)
|
||||
bar = Rectangle(width=w, height=bh,
|
||||
fill_color=interpolate_color(BLUE, RED, pct / 100),
|
||||
fill_opacity=0.85, stroke_color=LIGHT_GRAY, stroke_width=0.3)
|
||||
bar.move_to([lx + w / 2, y, 0])
|
||||
lbl = Text(f"{tok} {pct}%", font_size=10, color=WHITE)
|
||||
lbl.next_to(bar, RIGHT, buff=0.05)
|
||||
bars.add(bar); lbls.add(lbl)
|
||||
return VGroup(bars, lbls)
|
||||
|
||||
dists = [
|
||||
{"The": 72, "cat": 8, "sat": 6, "on": 4, "<unk>": 10},
|
||||
{"cat": 65, "sat": 12, "was": 8, "is": 6, "<unk>": 9},
|
||||
{"sat": 58, "slept": 15, "ran": 8, "jumped": 6, "<unk>": 13},
|
||||
{"on": 55, "down": 12, "quietly": 8, "and": 6, "<unk>": 19},
|
||||
{"the": 60, "a": 12, "top": 8, "floor": 5, "<unk>": 15},
|
||||
{"mat": 50, "table": 15, "chair": 8, "floor": 6, "<unk>": 21},
|
||||
]
|
||||
Y_DIST = -1.65
|
||||
|
||||
# ── Token sequence row ──
|
||||
SX = -4.0; Y_SEQ = 2.4; GAP = 0.70
|
||||
seq = VGroup()
|
||||
sos = tok_card("<s>", BLUE, BLUE).move_to([SX, Y_SEQ, 0])
|
||||
seq.add(sos)
|
||||
|
||||
# Position labels under each token
|
||||
pos_lbls = VGroup()
|
||||
pos0 = Text("0", font_size=9, color=DARK_GRAY).next_to(sos, DOWN, buff=0.08)
|
||||
pos_lbls.add(pos0)
|
||||
|
||||
# Step label (below everything)
|
||||
step_lbl = Text("Step 0 — [<s>] → ?", font_size=9, color=GRAY)
|
||||
step_lbl.move_to([0, -2.1, 0])
|
||||
|
||||
# ── Show static elements ──
|
||||
self.play(FadeIn(pipeline), FadeIn(arrows))
|
||||
self.play(FadeIn(layer_stack))
|
||||
self.play(FadeIn(kv_group))
|
||||
self.play(FadeIn(sos), Write(pos0))
|
||||
self.play(Write(step_lbl))
|
||||
self.wait(0.5)
|
||||
|
||||
# ── Generation loop ──
|
||||
for i, tok in enumerate(gen_tokens[1:], start=1):
|
||||
input_str = " ".join(gen_tokens[:i + 1])
|
||||
|
||||
# 1. Highlight last token (the input being processed)
|
||||
last = seq[-1]
|
||||
hl = SurroundingRectangle(last, color=YELLOW, stroke_width=2, buff=0.04)
|
||||
self.play(Create(hl), run_time=0.12)
|
||||
|
||||
# 2. Arrow from last token → Embedding
|
||||
in_arr = Arrow(last.get_bottom(), emb_node.get_top(), color=YELLOW,
|
||||
stroke_width=2, tip_length=0.08)
|
||||
self.play(FadeIn(in_arr, scale=0.5), run_time=0.1)
|
||||
|
||||
# 3. Cascade through pipeline
|
||||
self.play(
|
||||
*[p[0].animate.set_fill_opacity(0.6) for p in pipeline],
|
||||
run_time=0.12
|
||||
)
|
||||
self.play(
|
||||
*[p[0].animate.set_fill_opacity(0.25) for p in pipeline],
|
||||
run_time=0.1
|
||||
)
|
||||
self.play(FadeOut(in_arr), FadeOut(hl), run_time=0.08)
|
||||
|
||||
# 4. Show probability distribution
|
||||
dist_arr = Arrow(head.get_bottom(), head.get_bottom() + DOWN * 0.22,
|
||||
color=GRAY, stroke_width=1, tip_length=0.06)
|
||||
dist = build_dist(dists[i - 1], Y_DIST)
|
||||
self.play(FadeIn(dist_arr, scale=0.5), FadeIn(dist, scale=0.8), run_time=0.25)
|
||||
|
||||
# 5. Sampling: highlight top bar
|
||||
top_bar = dist[0][0]
|
||||
sample_hl = SurroundingRectangle(top_bar, color=YELLOW, stroke_width=1.5, buff=0.02)
|
||||
samp_lbl = Text("argmax", font_size=7, color=YELLOW)
|
||||
samp_lbl.next_to(dist, DOWN, buff=0.05)
|
||||
self.play(Create(sample_hl), Write(samp_lbl), run_time=0.2)
|
||||
self.wait(0.15)
|
||||
|
||||
# 6. Predicted token appears below distribution
|
||||
pred = tok_card(tok, YELLOW, YELLOW)
|
||||
target_x = SX + len(seq) * GAP
|
||||
pred.move_to([CX_BLK, Y_DIST - 0.45, 0])
|
||||
self.play(FadeIn(pred, scale=0.4), run_time=0.15)
|
||||
|
||||
# 7. Token rises to join sequence at top
|
||||
self.play(pred.animate.move_to([target_x, Y_SEQ, 0]), run_time=0.3)
|
||||
pos_lbl = Text(str(i), font_size=9, color=DARK_GRAY)
|
||||
pos_lbl.next_to(pred, DOWN, buff=0.08)
|
||||
self.play(
|
||||
pred[0].animate.set_fill(DARK_BLUE, 0.5).set_stroke(GRAY, 0.5),
|
||||
pred[1].animate.set_color(WHITE),
|
||||
Write(pos_lbl),
|
||||
FadeOut(sample_hl), FadeOut(samp_lbl),
|
||||
)
|
||||
seq.add(pred)
|
||||
pos_lbls.add(pos_lbl)
|
||||
|
||||
# 8. KV Cache: add K,V of this predicted token
|
||||
kv_x = KV_X + i * (kv_size + kv_gap)
|
||||
k_sq = Square(kv_size, fill_color=YELLOW, fill_opacity=0.4,
|
||||
stroke_color=YELLOW, stroke_width=0.5).move_to([kv_x, K_Y, 0])
|
||||
v_sq = Square(kv_size, fill_color=ORANGE, fill_opacity=0.4,
|
||||
stroke_color=ORANGE, stroke_width=0.5).move_to([kv_x, V_Y, 0])
|
||||
self.play(FadeIn(k_sq, scale=1.5), FadeIn(v_sq, scale=1.5), run_time=0.15)
|
||||
kv_k.add(k_sq); kv_v.add(v_sq)
|
||||
|
||||
# 9. Remove distribution and arrow
|
||||
self.play(FadeOut(dist), FadeOut(dist_arr), run_time=0.08)
|
||||
|
||||
# 10. Update step label
|
||||
new_lbl = Text(f"Step {i} — [{input_str}]", font_size=9, color=GRAY)
|
||||
new_lbl.move_to(step_lbl)
|
||||
self.play(Transform(step_lbl, new_lbl))
|
||||
self.wait(0.3)
|
||||
|
||||
self.wait(2.0)
|
||||
|
||||
|
||||
def orth_line(start, end, color=GRAY):
|
||||
|
||||
Reference in New Issue
Block a user