⚡ FlashDec

从 PagedAttention kernel 到可验证的单 GPU LLM decode runtime

A correctness-first decode runtime built with PyTorch, Triton and CUDA.

FlashDec 用一条可审计的数据路径研究单 token decode：请求如何拥有 paged KV block，调度器如何在容量压力下保证进展，多层 K/V 写入如何原子提交，以及 kernel 优化能否穿过 runtime 开销形成完整 step 收益。

项目面向单 GPU、caller-provided Q/K/V，不实现完整模型或 serving 服务。PyTorch reference 定义数值语义；PagedKVCache 独占 block、seq_len 与事务状态；Triton/CUDA 负责受这些状态约束的数据路径。

研究问题

问题	FlashDec 的回答	主要证据
logical token 如何映射到非连续 physical KV blocks？	独立 dense/paged reference、token-major cache、显式 block table 与 masked online softmax	warp selection · block-size · layout · default profile
谁拥有 block、seq_len 和 request lifecycle？	PagedKVCache 是唯一权威状态源；kernel 和 benchmark 不推进 lifecycle	Paged KV 设计 · DecodeEngine 设计
KV 容量不足时如何避免 boundary deadlock 与 starvation？	lifetime block commitment、FIFO + aging、公平 runnable subset、stale decision 拒绝	scheduler matrix
多层 token 如何只提交一次并支持整体回滚？	所有 layer 共享预留位置，按序写入，batch 原子 commit/abort	multi-layer matrix · transaction design
shared prefix 的收益究竟来自哪里？	只共享 immutable full blocks，以 refcount/LRU 管理；容量与 admission 收益和 latency 分开报告	8-trial confirmation
kernel 优化如何传播到系统，又如何公平比较外部实现？	分离 append-only、complete-step、profiler 和共同 kernel scope；保留负结果与不可比边界	性能报告 · FlashInfer baseline

完整的问题定义、假设和证据链见研究问题。

架构

RequestSpec / caller-provided Q,K,V
                 │
                 ▼
       Block-aware Scheduler
                 │ versioned decision
                 ▼
            DecodeEngine
                 │ token transaction
                 ▼
          PagedKVCache runtime
       ownership / refcount / rollback
                 │
       ┌─────────┴─────────┐
       ▼                   ▼
Fused RoPE + KV append   Triton paged decode
       └─────────┬─────────┘
                 ▼
      correctness + benchmark evidence

Scheduler 只产生版本化容量决策；DecodeEngine 编排执行；PagedKVCache 管理 ownership、事务、seq_len 和 shared-prefix lifetime。完整不变量见总体设计。

主要发现

• Scheduler 的主要价值是容量安全和进展保证：boundary case 中 lifetime policy 完成率为 100%，cancel baseline 为 50%，greedy baseline 为 0%。
• 75% shared-prefix hit 将 context physical blocks 从 64/64 降至 20/64，节省 68.8%（5.5 MiB）容量，并把固定池 admission 从 9/16 提高到 16/16；latency 没有稳定方向。
• Cache-owned trusted transaction 移除了每 layer 的重复 device reduction 与 host scalar sync；persistent-metadata 候选仅 13/16 组稳定，未采用。
• 固定 FlashInfer 0.6.15.post1 的共同 paged-decode kernel 对比中，FlashDec/FlashInfer p50 latency ratio 为 1.2003x（CUDA core）和 1.2284x（tensor core）；比值大于 1 表示 FlashInfer 更低延迟。该结论不外推到完整 runtime。
• 外部基线证据提交 d7d4feb 的 GPU full regression 为 453 passed, 94 subtests passed；GitHub Actions 运行不依赖 GPU 的仓库检查子集。测试计数绑定具体 commit 和环境，不作为滚动徽章。

图中数据由受版本控制的 Markdown summaries 派生。来源、trial 范围、ratio 方向和负结果见结果索引与数据快照。

快速开始

推荐 Linux/WSL 与 Python 3.10+。CUDA extension 需要与 PyTorch CUDA build 匹配的 Toolkit、NVCC 和 Ninja。

git clone https://github.com/cysong2025/flashdec.git
cd flashdec

python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip setuptools wheel
python -m pip install -e ".[dev,cuda-extension]"

python scripts/check_env.py
python -m pytest -q -ra

这条开发安装路径只适用于满足支持矩阵的环境，不构成对任意 Python/CUDA 组合的兼容性保证。请先核对 Python、PyTorch CUDA build、Toolkit/NVCC 与 GPU 架构；已验证组合和分层命令见兼容性说明与复现指南。FlashInfer 对比必须使用独立环境和 constraints/flashinfer-cu128.txt。

无需 GPU 的仓库检查：

python scripts/check_docs.py
python scripts/check_release.py --require-evidence --require-public
python -m compileall -q flashdec tests benchmarks scripts

API 速览

Paged decode：

import flashdec

head_dim = q.shape[-1]
out = flashdec.decode(
    q,
    k_cache,
    v_cache,
    block_tables,
    seq_lens,
    sm_scale=head_dim**-0.5,
    block_size=32,
)

Multi-layer token transaction：

tx = engine.begin_step(request_ids)
for layer_idx, (q, k, v) in enumerate(zip(q_by_layer, k_by_layer, v_by_layer)):
    engine.step_layer(tx, layer_idx, q, k, v)
result = engine.commit_step(tx)

任一 step_layer() 失败都会 abort 整个 token；调用方必须丢弃此前 layer outputs。接口契约见API 文档。

仓库结构

路径	内容
flashdec/	Paged KV runtime、scheduler、DecodeEngine 与 Triton/CUDA data path
tests/	reference、状态机、kernel、失败路径与证据校验
benchmarks/	benchmark runners、profilers 和 strict summarizers
benchmarks/results/	审核后提交的正式 Markdown 证据；原始 CSV/log 不进入 Git
docs/	研究问题、设计、兼容性、性能方法和复现说明
scripts/	环境、文档、证据与仓库一致性检查

推荐阅读顺序：

1. 研究问题
2. 总体设计
3. API 文档
4. 性能报告
5. 结果索引
6. 复现指南

完整导航见文档索引。

范围边界

FlashDec 包含单 GPU single-token decode、FP16/BF16、MHA/GQA/MQA、caller-provided Q/K/V、transactional context import 和 caller-built immutable shared-prefix blocks。

FlashDec 不包含完整 Transformer/prefill forward、tokenizer、sampling、HTTP/RPC serving、生产级 continuous batching、TP/PP、多机、swap/offload 或自动 prefix content hashing。不同 shape、layout、scheduler 或计时边界的结果不能直接解释为工业 serving 系统的 speedup。

贡献前请阅读贡献指南与行为准则；使用问题见支持说明，安全问题见安全政策，引用信息见CITATION.cff。

License

FlashDec 采用 Apache License 2.0 开源。