rednote-machine-learning/RedKnot

## 关于 **RedKnot** 是一个建立在 [SGLang](https://github.com/sgl-project/sglang) 基础上的长上下文推理加速集成。其核心理念是：**并非所有的 attention head 都需要完整的 KV，也并非所有的 token 都需要经过完整的 FFN**。RedKnot 通过以下方式实现这一点： - **Head 分类**：每个 `(layer, kv_head)` 被划分为四类之一 —— `global / local / retrieval / dense` —— 并根据类别决定 KV 存储和复用策略； - **离线 KV 复用 + RoPE 重定位**：可复用片段的 KV 被离线存储；在服务时仅选择性地重新计算必要的 token，并且 RoPE 重定位可保证数值对齐； - **弹性稀疏 / 稀疏 FFN**：基于 attention 重要性的 token 选择性 FFN，跳过低贡献 token 的前馈计算； - **SegPagedAttention 运行时**：基于 head 的 page table + 分段 KV 存储，允许不同的 head 类别拥有不同的可见窗口； - **DeepSeek-V4** 和**完整的 Qwen 3.5 系列**将在下一个版本中完全更新和适配；目前仅开源基础版本。 在保持**近乎无损的准确率**（在某些场景下甚至优于 dense baseline）的同时，它将长上下文 prefill 的 **FLOPs 减少了大约 50%–70%**，并带来 **1.35 倍至 2.2 倍的 TTFT 加速**（随着上下文长度的增加，收益会更大）。 - **论文**：**RedKnot: Efficient Long-Context LLM Serving with Head-Aware KV Reuse and SegPagedAttention** — Yang Liu, ZhaoKai Luo, HuaYi Jin, ZhiYong Wang, RuoZhou He, BoYu Wang, Guanjie Chen, Junhao Hao () ## 核心思想 | 机制 | 描述 | 代码位置 | |---|---|---| | Head 分类配置 | `global / local / retrieval / dense` 四类策略 + JSON 加载 | `redknot/head_config.py`, `head_profiler.py` | | 离线 KV cache + RoPE 重定位 | 片段级离线 KV 存储和重定位（数值对齐） | `redknot/offline_cache.py`, `rope_helper.py` | | Head 感知的 attention 恢复 | FlashAttention-2 / FA-3 分桶 attention | `redknot/ops_flash.py`, `ops_flash3.py` | | 稀疏 FFN（弹性稀疏） | Token 选择性 FFN，根据重要性跳过计算 | `redknot/sparse_ffn.py` | | SegPagedAttention 运行时 | 基于 head 的 page table + 分段 KV 存储 | `redknot/segpaged.py`, `segpaged_v2/` | | DeepSeek-V4 MLA 集成 | 复用 indexer top-k 以进行选择性重计算 | `redknot/deepseek_v4_mla.py`, `dsv4_offline_reuse.py` | | PD KV 传输 / Head 感知调度 | Head 类 KV 分片传输和容量模型 | `redknot/pd_transfer.py`, `scheduler.py` | 有关更详细的阶段计划，请参见 `python/sglang/srt/layers/attention/redknot/ROADMAP.md`。 ## 快速开始 RedKnot 复用了 SGLang 的安装流程： ``` # 安装（开发模式） pip install -e "python[all]" ``` 某些模型（`Qwen3.5-*` 的 `qwen3_5_moe` 架构，以及 `DeepSeek-V4`）需要 transformers 5.x。 代码库自带 `.venv_tf5` (transformers 5.12.0)。系统的 transformers 4.57 无法加载这些模型。 - 安装： - 快速开始： ## 基准测试（RAG 准确率 / 速度） 所有基准测试脚本均位于 `test/srt/redknot/`；相关的脚本、图表和文档归档在 `test/srt/redknot/utils/` 下。 每个基准测试都使用一个**诚实的 dense baseline**（一次完整的 FlashAttention-2 prefill）来与 RedKnot 的 head 类 KV 复用路径进行对比，报告准确率（SQuAD F1 / EM）、速度（TTFT、加速比、decode tok/s）以及后台开销。 ### Qwen3-32B — HotpotQA ✅ 最佳 | 上下文 | 基准 F1 | RedKnot F1 | 基准 TTFT | RedKnot TTFT | 加速比 | 节省的 FLOPs | |---|---|---|---|---|---|---| | 16K | 0.750 | **1.000** | 3.24s | 2.33s | **1.39x** | 69.2% | | 24K | 1.000 | **1.000** | 5.25s | 2.96s | **1.77x** | 70.9% | | 32K | 0.750 | **1.000** | 7.74s | 4.02s | **1.93x** | 72.2% | RedKnot 的 F1 **始终 ≥ 基准**（无损或更好），TTFT 加速比随上下文增加而增长（1.39→1.93x），节省的 FLOPs 稳定在 ~70%。 ### Qwen3.5-35B-A3B (MoE) — LongBench | 上下文 | 数据集 | 标准 F1 | RedKnot F1 | 节省的计算量 | TTFT 加速比 | |---|---|---|---|---|---| | 16K | triviaqa | 1.000 | 1.000 | 46.4% | **1.87x** | | 32K | multifieldqa_en | 0.792 | 0.576 | 50.4% | **2.02x** | | 64K | triviaqa | 0.875 | 0.750 | 53.8% | **2.16x** | TTFT 加速比随上下文增加而增长（1.87→2.16x）；在 16K 时无损，但在长上下文下存在一些性能下降（这是线性 attention + MoE 稀疏带来的代价）。 ### Mistral-7B-Instruct-v0.3 — HotpotQA | 上下文 | 基准 F1 | RedKnot F1 | 基准 TTFT | RedKnot TTFT | 加速比 | 节省的 FLOPs | |---|---|---|---|---|---|---| | 16K | 0.250 | 0.475 | 0.70s | 0.52s | **1.35x** | 51.5% | | 24K | 0.250 | 0.250 | 1.12s | 0.80s | **1.39x** | 50.2% | | 32K | 0.688 | 0.100 | 1.67s | 1.24s | **1.35x** | 49.1% | 对于这个较小的 7B 模型，基准模型本身在长上下文 RAG 方面表现较弱，且 F1 波动较大；系统指标（TTFT 加速比约 1.35 倍 / 节省的 FLOPs 约 50%）保持稳定。 ### 已知问题 - **Llama-3.3-70B-Instruct**：基准测试正常运行，但在长上下文 LongBench 下，RedKnot 的 decode 路径出现性能下降（重复 token），并且单 GPU INT4 很容易 OOM，而多 GPU bf16 会引发跨设备错误。这是 RedKnot 在 Llama3.3 上预先存在的算法/配置问题，有待单独调查 `driver_batched` 的 Llama 兼容性以及 `head_class/llama-70B_*.json` 配置的质量。 ## 如何运行 基准测试脚本均位于 `test/srt/redknot/` 中。在运行时，它们依赖于同一文件夹下的 `head_class/`、`sparse_ffn_params/` 和 `datasets/` 配置/数据目录，以及 `utils/`（例如 `fp8_offline_patch.py`）。 ### 一键复现 ``` cd test/srt/redknot # 默认 small/medium 模型（HotpotQA 上的 Mistral/Qwen3 + LongBench 上的 Llama/Qwen35） bash run_all_rag.sh # 自定义模型和尺寸 RK_MODELS="mistral qwen3" RK_SAMPLES=4 RK_LENGTHS=16K,24K,32K \ bash run_all_rag.sh ``` ### 运行所有 RAG 基准测试 依次运行五个模型（Qwen3.5-MoE / Qwen3 / Mistral / Llama3.3 / DeepSeek-V4）的 RAG 基准测试： ``` cd test/srt/redknot python benchmark_RedKnot_Qwen35_RAG.py python benchmark_RedKnot_Qwen3_RAG.py python benchmark_RedKnot_Mistral_RAG.py python benchmark_RedKnot_Llama3.3_RAG.py python benchmark_RedKnot_DeepSeekV4_RAG.py ``` ### 最小化单模型检查（几分钟内出结果） ``` # Qwen3-32B (INT4 NF4, single GPU) REDKNOT_N_SAMPLES=1 REDKNOT_LENGTHS=16K REDKNOT_MAX_NEW=8 \ CUDA_VISIBLE_DEVICES=0 python test/srt/redknot/benchmark_RedKnot_Qwen3_RAG.py # Qwen3.5-35B-A3B / Qwen3.5-397B-A17B (MoE, requires transformers 5, use .venv_tf5) PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True HF_HUB_OFFLINE=1 \ REDKNOT_N_SAMPLES=1 REDKNOT_MAX_NEW=8 CUDA_VISIBLE_DEVICES=0,1 \ .venv_tf5/bin/python test/srt/redknot/benchmark_RedKnot_Qwen35_397B_RAG.py # Mistral-7B-Instruct-v0.3 (bf16, single GPU) REDKNOT_N_SAMPLES=4 CUDA_VISIBLE_DEVICES=0 \ python test/srt/redknot/benchmark_RedKnot_Mistral_RAG.py # Llama-3.3-70B-Instruct (INT4 NF4, single GPU) REDKNOT_N_SAMPLES=3 CUDA_VISIBLE_DEVICES=0 \ python test/srt/redknot/benchmark_RedKnot_Llama3.3_RAG.py # DeepSeek-V4 (MLA + indexer, requires large VRAM / .venv_tf5) CUDA_VISIBLE_DEVICES=0 \ .venv_tf5/bin/python test/srt/redknot/benchmark_RedKnot_DeepSeekV4_RAG.py ``` ### 关键环境变量 | 变量 | 描述 | |---|---| | `REDKNOT_N_SAMPLES` | 评估样本数 | | `REDKNOT_LENGTHS` | 上下文长度（HotpotQA 模型，例如 `16K,24K,32K`） | | `REDKNOT_DATASETS` | LongBench 数据集（LongBench 模型） | | `REDKNOT_MAX_NEW` | 最大生成 token 数 | | `REDKNOT_DTYPE` | `int4` 或 `bf16` | | `REDKNOT_COMPILE` | 是否启用 `torch.compile` (`0`/`1`) | | `CUDA_VISIBLE_DEVICES` | 可见的 GPU | 所有运行日志都保存在 `test/srt/redknot/rag_logs/` 下。 ## 致谢 RedKnot 构建在 [SGLang](https://github.com/sgl-project/sglang) 之上，并复用了其生态系统中许多项目的设计和实现： - [SGLang](https://github.com/sgl-project/sglang) - [vLLM](https://github.com/vllm-project/vllm)

标签：DLL 劫持, KV缓存复用, SGLang, Vectored Exception Handling, 凭据扫描, 大语言模型, 推理优化, 模型推理加速, 稀疏注意力, 逆向工具, 长上下文处理